diff --git a/.env.example b/.env.example index 61cb94d..bfa9e7c 100644 --- a/.env.example +++ b/.env.example @@ -10,6 +10,10 @@ COMPUTE_ID=compute-env-id WORK_SPACE=workspace-id WORK_DIR=s3://bucket/workdir +# NCI Gadi project code used for PBS submissions (clusterOptions) across all +# Nextflow workflows. Optional; defaults to "yz52" if unset. +GADI_PROJECT=yz52 + # AWS S3 Configuration for file uploads AWS_ACCESS_KEY_ID=your-access-key-id AWS_SECRET_ACCESS_KEY=your-secret-access-key diff --git a/.github/workflows/check-migrations.yml b/.github/workflows/check-migrations.yml index bb1ca37..29964d9 100644 --- a/.github/workflows/check-migrations.yml +++ b/.github/workflows/check-migrations.yml @@ -2,7 +2,7 @@ name: Check Alembic Migrations on: pull_request: - branches: [main] + branches: [dev] paths: - "alembic/**" - "alembic.ini" diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index ac039d5..4cb292b 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -2,9 +2,9 @@ name: Lint on: push: - branches: [main] + branches: [dev] pull_request: - branches: [main] + branches: [dev] jobs: ruff: diff --git a/.github/workflows/test-coverage.yml b/.github/workflows/test-coverage.yml index 3632112..00dcc73 100644 --- a/.github/workflows/test-coverage.yml +++ b/.github/workflows/test-coverage.yml @@ -2,9 +2,9 @@ name: Coverage on: push: - branches: ["main"] + branches: ["dev"] pull_request: - branches: ["main"] + branches: ["dev"] jobs: tests: diff --git a/alembic/versions/20260717_103718_add_tool_and_require_core_fields_on__7a377a05de6e.py b/alembic/versions/20260717_103718_add_tool_and_require_core_fields_on__7a377a05de6e.py new file mode 100644 index 0000000..ae06edb --- /dev/null +++ b/alembic/versions/20260717_103718_add_tool_and_require_core_fields_on__7a377a05de6e.py @@ -0,0 +1,41 @@ +"""add_tool_and_require_core_fields_on_workflows""" + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '7a377a05de6e' +down_revision = '74ab4858580d' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('workflows', sa.Column('tool', sa.Text(), nullable=True)) + op.alter_column('workflows', 'repo_url', + existing_type=sa.TEXT(), + nullable=False) + op.alter_column('workflows', 'default_revision', + existing_type=sa.TEXT(), + nullable=False) + op.alter_column('workflows', 'config_path', + existing_type=sa.TEXT(), + nullable=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column('workflows', 'config_path', + existing_type=sa.TEXT(), + nullable=True) + op.alter_column('workflows', 'default_revision', + existing_type=sa.TEXT(), + nullable=True) + op.alter_column('workflows', 'repo_url', + existing_type=sa.TEXT(), + nullable=True) + op.drop_column('workflows', 'tool') + # ### end Alembic commands ### diff --git a/alembic/versions/20260717_105701_add_unique_constraints_for_workflow__14cff803a3b9.py b/alembic/versions/20260717_105701_add_unique_constraints_for_workflow__14cff803a3b9.py new file mode 100644 index 0000000..5f05e41 --- /dev/null +++ b/alembic/versions/20260717_105701_add_unique_constraints_for_workflow__14cff803a3b9.py @@ -0,0 +1,25 @@ +"""add_unique_constraints_for_workflow_name_tool""" + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '14cff803a3b9' +down_revision = '7616ae331fe5' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_unique_constraint(op.f('uq_workflows_name'), 'workflows', ['name', 'tool']) + op.create_index('uq_workflows_name_null_tool', 'workflows', ['name'], unique=True, postgresql_where=sa.text('tool IS NULL')) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index('uq_workflows_name_null_tool', table_name='workflows', postgresql_where=sa.text('tool IS NULL')) + op.drop_constraint(op.f('uq_workflows_name'), 'workflows', type_='unique') + # ### end Alembic commands ### diff --git a/alembic/versions/20260717_153902_drop_stale_name_only_unique_index_on__7616ae331fe5.py b/alembic/versions/20260717_153902_drop_stale_name_only_unique_index_on__7616ae331fe5.py new file mode 100644 index 0000000..86872cb --- /dev/null +++ b/alembic/versions/20260717_153902_drop_stale_name_only_unique_index_on__7616ae331fe5.py @@ -0,0 +1,24 @@ +"""drop_stale_name_only_unique_index_on_workflows""" + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '7616ae331fe5' +down_revision = '7a377a05de6e' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Some environments have a plain unique index named 'uq_workflows_name' on + # workflows.name alone, predating this feature and untracked by any prior + # migration (likely created out-of-band). It collides with the + # same-named (name, tool) constraint the next migration creates. + op.execute("DROP INDEX IF EXISTS uq_workflows_name") + + +def downgrade() -> None: + # The dropped index was untracked drift, not a schema feature; nothing to restore. + pass diff --git a/alembic/versions/20260724_094600_system_status_incidents_a6ee1af2c156.py b/alembic/versions/20260724_094600_system_status_incidents_a6ee1af2c156.py new file mode 100644 index 0000000..2d559ce --- /dev/null +++ b/alembic/versions/20260724_094600_system_status_incidents_a6ee1af2c156.py @@ -0,0 +1,33 @@ +"""system_status_incidents""" + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = 'a6ee1af2c156' +down_revision = '14cff803a3b9' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('system_status_incidents', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('component', sa.Text(), nullable=False), + sa.Column('status', sa.Text(), nullable=False), + sa.Column('started_at', sa.DateTime(timezone=True), nullable=False), + sa.Column('ended_at', sa.DateTime(timezone=True), nullable=True), + sa.Column('message', sa.Text(), nullable=True), + sa.PrimaryKeyConstraint('id', name=op.f('pk_system_status_incidents')) + ) + op.create_index('ix_system_status_incidents_component_started_at', 'system_status_incidents', ['component', 'started_at'], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index('ix_system_status_incidents_component_started_at', table_name='system_status_incidents') + op.drop_table('system_status_incidents') + # ### end Alembic commands ### diff --git a/alembic/versions/20260728_145217_service_usage_4f26ab5038ec.py b/alembic/versions/20260728_145217_service_usage_4f26ab5038ec.py new file mode 100644 index 0000000..8935f18 --- /dev/null +++ b/alembic/versions/20260728_145217_service_usage_4f26ab5038ec.py @@ -0,0 +1,23 @@ +"""service_usage""" + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '4f26ab5038ec' +down_revision = 'a6ee1af2c156' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('workflow_runs', sa.Column('service_usage', sa.Float(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column('workflow_runs', 'service_usage') + # ### end Alembic commands ### diff --git a/app/db/admin.py b/app/db/admin.py index 184704d..c5e26db 100644 --- a/app/db/admin.py +++ b/app/db/admin.py @@ -126,6 +126,7 @@ class WorkflowAdmin(ModelView): fields = [ "id", "name", + "tool", "description", "repo_url", "default_revision", @@ -135,6 +136,7 @@ class WorkflowAdmin(ModelView): _NULLABLE_FIELDS = ( "description", + "tool", "prerun_script_path", ) @@ -158,20 +160,22 @@ async def repr(self, obj: Any, request: Request) -> str: class WorkflowRunAdmin(ModelView): fields = [ - "id", + "submission_timestamp", + HasOne("workflow", identity="workflow"), "workflow_id", "tool", - "owner_user_id", - HasOne("workflow", identity="workflow"), HasOne("owner", identity="app-user"), + "owner_user_id", "seqera_run_id", "run_name", - "binder_name", + "service_usage", JSONField("submitted_form_data"), + "binder_name", "work_dir", - "submission_timestamp", + "id", ] - exclude_fields_from_list = "submitted_form_data" + exclude_fields_from_list = ["submitted_form_data"] + fields_default_sort = [("submission_timestamp", True)] async def repr(self, obj: Any, request: Request) -> str: return f"{obj.run_name}" diff --git a/app/db/models/core.py b/app/db/models/core.py index 31ded49..ac5736f 100644 --- a/app/db/models/core.py +++ b/app/db/models/core.py @@ -7,11 +7,14 @@ JSON, BigInteger, DateTime, + Float, ForeignKey, + Index, Numeric, PrimaryKeyConstraint, Text, UniqueConstraint, + text, ) from sqlalchemy.dialects.postgresql import INET, UUID from sqlalchemy.orm import Mapped, Session, mapped_column, relationship @@ -39,14 +42,32 @@ class AppUser(Base): class Workflow(Base): __tablename__ = "workflows" + __table_args__ = ( + # At most one row per (name, tool) for workflows with tool-specific + # configs (e.g. de-novo-design: bindcraft vs rfdiffusion). + UniqueConstraint("name", "tool"), + # At most one generic (tool-independent) row per name. sqlite_where is + # needed alongside postgresql_where so the tests' SQLite schema (built + # via Base.metadata.create_all) enforces the same partial uniqueness + # as the Postgres migration; without it, SQLite silently drops the + # WHERE clause and the index becomes a full unique constraint on name. + Index( + "uq_workflows_name_null_tool", + "name", + unique=True, + postgresql_where=text("tool IS NULL"), + sqlite_where=text("tool IS NULL"), + ), + ) id: Mapped[UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid4) name: Mapped[str] = mapped_column(Text, nullable=False) description: Mapped[str | None] = mapped_column(Text, nullable=True) - repo_url: Mapped[str | None] = mapped_column(Text, nullable=True) - default_revision: Mapped[str | None] = mapped_column(Text, nullable=True) - config_path: Mapped[str | None] = mapped_column(Text, nullable=True) + repo_url: Mapped[str] = mapped_column(Text, nullable=False) + default_revision: Mapped[str] = mapped_column(Text, nullable=False) + config_path: Mapped[str] = mapped_column(Text, nullable=False) prerun_script_path: Mapped[str | None] = mapped_column(Text, nullable=True) + tool: Mapped[str | None] = mapped_column(Text, nullable=True) runs: Mapped[list[WorkflowRun]] = relationship(back_populates="workflow") @@ -72,6 +93,7 @@ class WorkflowRun(Base): DateTime(timezone=True), nullable=True ) tool: Mapped[str | None] = mapped_column(Text, nullable=True) + service_usage: Mapped[float | None] = mapped_column(Float, nullable=True) owner: Mapped[AppUser] = relationship(back_populates="workflow_runs") workflow: Mapped[Workflow | None] = relationship(back_populates="runs") diff --git a/app/db/models/system_status.py b/app/db/models/system_status.py index b01b9bd..be3aff2 100644 --- a/app/db/models/system_status.py +++ b/app/db/models/system_status.py @@ -1,9 +1,9 @@ -"""Database-backed cache for runtime system health.""" +"""Database-backed cache and history for runtime system health.""" from datetime import UTC, datetime from typing import Any -from sqlalchemy import JSON, DateTime, Text +from sqlalchemy import JSON, DateTime, Index, Integer, Text from sqlalchemy.orm import Mapped, mapped_column from ...schemas.health import SystemStatus @@ -38,3 +38,24 @@ def get_status(self) -> SystemStatus: Return a SystemStatus object from the cache payload. """ return SystemStatus.model_validate(self.payload) + + +class SystemStatusIncident(Base): + """ + One row per downtime period for a monitored component: opened when a probe + first reports ``degraded``/``unhealthy``, closed when it next reports + ``healthy``. Healthy checks are never recorded, so storage scales with the + number of actual outages rather than with polling frequency. + """ + + __tablename__ = "system_status_incidents" + __table_args__ = ( + Index("ix_system_status_incidents_component_started_at", "component", "started_at"), + ) + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + component: Mapped[str] = mapped_column(Text, nullable=False) + status: Mapped[str] = mapped_column(Text, nullable=False) + started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) + ended_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + message: Mapped[str | None] = mapped_column(Text, nullable=True) diff --git a/app/routes/system_status.py b/app/routes/system_status.py index 2fd0fcc..fc6b9d0 100644 --- a/app/routes/system_status.py +++ b/app/routes/system_status.py @@ -12,11 +12,13 @@ from __future__ import annotations +from datetime import UTC, datetime, timedelta + from fastapi import APIRouter, Depends, Query from sqlalchemy.orm import Session from ..db.admin import require_admin_access -from ..schemas.health import SystemStatusAdminResponse +from ..schemas.health import SystemStatusAdminResponse, SystemStatusDowntimeResponse from ..services import health from .dependencies import get_db @@ -25,6 +27,9 @@ dependencies=[Depends(require_admin_access)], ) +# Bounds how far back /system-status/history can be queried. +_MAX_HISTORY_HOURS = 24 * 30 + @router.get("/system-status", response_model=SystemStatusAdminResponse) async def get_admin_system_status( @@ -37,3 +42,24 @@ async def get_admin_system_status( """Return verbose, admin-only runtime health of the submission components.""" status_obj = await health.get_system_status(db, force_refresh=refresh) return SystemStatusAdminResponse.model_validate(health.to_admin_dict(status_obj)) + + +@router.get("/system-status/history", response_model=SystemStatusDowntimeResponse) +async def get_admin_system_status_history( + hours: int = Query( + default=24, + ge=1, + le=_MAX_HISTORY_HOURS, + description="How many hours of downtime history to return", + ), + component: str | None = Query(default=None, description="Restrict to a single component name"), + db: Session = Depends(get_db), +) -> SystemStatusDowntimeResponse: + """Return per-component downtime incidents, so admins can see how frequently + and for how long a component has been degraded/unhealthy over time.""" + until = datetime.now(UTC) + since = until - timedelta(hours=hours) + incidents = health.get_incidents(db, since=since, until=until, component=component) + return SystemStatusDowntimeResponse.model_validate( + health.to_downtime_dict(incidents, since=since, until=until, hours=hours) + ) diff --git a/app/routes/workflow/jobs.py b/app/routes/workflow/jobs.py index 8824fd6..40bf523 100644 --- a/app/routes/workflow/jobs.py +++ b/app/routes/workflow/jobs.py @@ -33,6 +33,7 @@ get_owned_run_by_id, get_user_job_list_rows, parse_submit_datetime, + sync_service_usage, ) from ...services.seqera import describe_workflow from ...services.seqera_client import cancel_workflow_raw, delete_workflow_raw, delete_workflows_raw @@ -200,6 +201,10 @@ async def _fetch_seqera(user_run: UserJobListRow) -> tuple[str, dict[str, object if score is None: score = await ensure_completed_run_score(db, owned_run, ui_status) + # Also sync service usage for runs, if needed + if owned_run.service_usage is None: + await sync_service_usage(db, run=owned_run, ui_status=ui_status) + jobs.append( JobListItem( id=run_id, @@ -259,6 +264,9 @@ async def get_job_details( score = await ensure_completed_run_score(db, owned_run, ui_status) if ui_status != "Completed": score = None + # Sync service usage - not returned to user but currently this is where + # we trigger syncs + await sync_service_usage(db, run=owned_run, ui_status=ui_status) raw_tool: str | None = getattr(owned_run, "tool", None) or None if not raw_tool: diff --git a/app/routes/workflows.py b/app/routes/workflows.py index 5bf2e54..f1bf888 100644 --- a/app/routes/workflows.py +++ b/app/routes/workflows.py @@ -12,9 +12,8 @@ from fastapi import APIRouter, Depends, HTTPException, Query, status from pydantic import ValidationError -from sqlalchemy import CursorResult, func, select, update +from sqlalchemy import CursorResult, func, or_, select, update from sqlalchemy.orm import Session -from unidecode import unidecode from ..db.models import QueuedJob from ..db.models.core import AppUser, RunInput, RunMetric, S3Object, Workflow, WorkflowRun @@ -25,12 +24,14 @@ ListRunsResponse, RunInputPresignedUrlResponse, S3DatasetUploadResponse, + SinglePredictionEntity, WispsDatasetUploadRequest, WispsFormData, WorkflowFormData, WorkflowLaunchPayload, WorkflowLaunchResponse, WorkflowUserDetails, + validate_single_prediction_entities, ) from ..services.bindflow_executor import _get_required_env, prepare_bindflow_workflow from ..services.credits import ( @@ -45,6 +46,7 @@ upload_csv_to_s3, upload_wisps_samplesheet_to_s3, ) +from ..services.proteindj_executor import prepare_proteindj_workflow from ..services.proteinfold_executor import prepare_proteinfold_workflow from ..services.s3 import S3ConfigurationError, S3ServiceError, generate_presigned_url from ..services.seqera_errors import SeqeraConfigurationError @@ -137,6 +139,43 @@ def _extract_final_design_count(form_data: WorkflowFormData | None) -> int | Non return parsed if parsed >= 1 else None +def _validate_single_prediction_form(form_data: WorkflowFormData, tool: str) -> None: + """Validate single-prediction entity limits; raise HTTPException on failure. + + Guards against jobs with too many entities or has no protein input + """ + raw_entities = form_data.extra_fields.get("entities") + if not isinstance(raw_entities, list): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="'entities' is required in formData for single-prediction.", + ) + try: + entities = [SinglePredictionEntity.model_validate(item) for item in raw_entities] + except ValidationError as exc: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="Invalid entity data in formData for single-prediction.", + ) from exc + + raw_potentials = form_data.extra_fields.get("boltz_use_potentials") + boltz_use_potentials = ( + raw_potentials.strip().lower() in ("true", "1", "yes") + if isinstance(raw_potentials, str) + else bool(raw_potentials) + ) + + try: + validate_single_prediction_entities( + entities, tool, boltz_use_potentials=boltz_use_potentials + ) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=str(exc), + ) from exc + + @router.post("/me/sync") async def sync_current_user( current_user_id: UUID = Depends(get_current_user_id), @@ -208,16 +247,27 @@ async def launch_workflow( ) # Workflow repo_url and revision come from the DB entry for this workflow name - # ("single-prediction", "de-novo-design", etc.). + # ("single-prediction", "de-novo-design", etc.). A workflow's tool column is + # NULL for a single row shared by all of its tools (e.g. single-prediction), + # or set on multiple rows when each tool needs its own repo_url/config_path/ + # default_revision (e.g. de-novo-design: bindcraft vs rfdiffusion). An exact + # tool match is preferred over the generic NULL-tool row when both exist. + tool_matches = func.lower(Workflow.tool) == selected_tool.lower() workflow = db_session.scalar( - select(Workflow).where(func.lower(Workflow.name) == requested_workflow) + select(Workflow) + .where( + func.lower(Workflow.name) == requested_workflow, + or_(Workflow.tool.is_(None), tool_matches), + ) + .order_by(tool_matches.desc()) + .limit(1) ) if not workflow: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=( - f"Workflow '{payload.launch.workflow}' is not configured in workflows table. " - "Seed the workflows catalog before launching." + f"Workflow '{payload.launch.workflow}' with tool '{selected_tool}' is not " + "configured in workflows table. Seed the workflows catalog before launching." ), ) @@ -234,28 +284,16 @@ async def launch_workflow( ) user = db_session.execute( - select(AppUser.email, AppUser.name).where(AppUser.id == current_user_id) + select(AppUser.email).where(AppUser.id == current_user_id) ).one_or_none() if not user or not user.email: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Could not retrieve user details required for workflow launch.", ) - user_email = user.email - # removes everything that isn't a letter, digit, or space - name = unidecode(user.name or "") - full_name = re.sub(r"[^a-zA-Z0-9 ]", "", name).replace(" ", "_") - institute = user_email.split("@")[-1] if "@" in user_email else None - ip_address: str | None = launch_ip or None - - full_name = _require_launch_var("full_name", full_name) - institute = _require_launch_var("institute", institute) - ip_address = _require_launch_var("ip_address", ip_address) user_details = WorkflowUserDetails( - user_email=user_email, - full_name=full_name, - institute=institute, - ip_address=ip_address, + user_email=user.email, + ip_address=_require_launch_var("ip_address", launch_ip or None), ) # Authoritative credit cost (server-side, non-spoofable). Only charged for @@ -330,6 +368,9 @@ async def launch_workflow( detail=f"'{missing}' is required in formData for {workflow_name}.", ) from exc + if workflow_name in ("single-prediction", "proteinfold"): + _validate_single_prediction_form(payload.formData, selected_tool) + try: queued_job: QueuedJob seqera_run_name = build_unique_run_name(payload.launch.runName or "") @@ -352,23 +393,36 @@ async def launch_workflow( user_details=user_details, ) elif workflow_name in ("de-novo-design", "bindflow", "bindcraft"): - # de-novo-design → bindflow executor. - # selected_tool carries the chosen algorithm ("bindcraft", "rfdiffusion"). + # de-novo-design → bindflow executor (bindcraft) or proteindj executor + # (rfdiffusion), depending on the chosen algorithm. tool_mode = selected_tool - bindcraft_launch_form = payload.launch.model_copy(update={"runName": seqera_run_name}) - queued_job = await prepare_bindflow_workflow( - bindcraft_launch_form, - s3_input_key, - db_session=db_session, - workflow_run=workflow_run, - pipeline=workflow.repo_url, - config_path=workflow.config_path, - revision=workflow.default_revision, - output_id=str(run_id), - mode=tool_mode, - form_data=payload.formData, - user_details=user_details, - ) + de_novo_launch_form = payload.launch.model_copy(update={"runName": seqera_run_name}) + if tool_mode.lower() == "rfdiffusion": + queued_job = await prepare_proteindj_workflow( + de_novo_launch_form, + db_session=db_session, + workflow_run=workflow_run, + pipeline=workflow.repo_url, + config_path=workflow.config_path, + revision=workflow.default_revision, + output_id=str(run_id), + form_data=payload.formData, + user_details=user_details, + ) + else: + queued_job = await prepare_bindflow_workflow( + de_novo_launch_form, + s3_input_key, + db_session=db_session, + workflow_run=workflow_run, + pipeline=workflow.repo_url, + config_path=workflow.config_path, + revision=workflow.default_revision, + output_id=str(run_id), + mode=tool_mode, + form_data=payload.formData, + user_details=user_details, + ) elif workflow_name in ("interaction-screening", "bulk-prediction"): assert wisps_form_data is not None wisps_launch_form = payload.launch.model_copy(update={"runName": seqera_run_name}) @@ -406,7 +460,7 @@ async def launch_workflow( .values( credit=AppUser.credit - run_credit_cost, credit_updated_at=datetime.now(UTC), - credit_updated_by=user_email, + credit_updated_by=user_details.user_email, ) ), ) diff --git a/app/run_scheduler.py b/app/run_scheduler.py index 5f1779e..bd0d609 100644 --- a/app/run_scheduler.py +++ b/app/run_scheduler.py @@ -3,10 +3,17 @@ import typer from apscheduler.triggers.cron import CronTrigger from apscheduler.triggers.interval import IntervalTrigger +from dotenv import load_dotenv from loguru import logger -from app.scheduler import SCHEDULER -from app.scheduler.jobs import refresh_user_credits, submit_pending_jobs +# Load .env before importing scheduler modules, which read required env vars (e.g. +# SEQERA_API_URL) at call time. Unlike app/main.py, this entrypoint runs standalone +# (`python -m app.run_scheduler`) and never imports main.py, so without this call +# .env values are never loaded into the process environment. +load_dotenv() + +from app.scheduler import SCHEDULER # noqa: E402 +from app.scheduler.jobs import refresh_user_credits, submit_pending_jobs # noqa: E402 SUBMIT_INTERVAL = IntervalTrigger(minutes=5) MONTHLY_TRIGGER = CronTrigger(day=1, hour=1, minute=0, timezone="UTC") diff --git a/app/scheduler/jobs.py b/app/scheduler/jobs.py index 13895d8..3ce5ec2 100644 --- a/app/scheduler/jobs.py +++ b/app/scheduler/jobs.py @@ -1,4 +1,5 @@ import asyncio +import os from collections.abc import Awaitable from datetime import UTC, datetime, timedelta from typing import Protocol, cast @@ -11,16 +12,23 @@ from ..db.models.job_queue import QueuedJob from ..routes.dependencies import get_db from ..schemas.workflows import WorkflowName -from ..services import health +from ..services import health, seqera from ..services.bindflow_executor import launch_bindflow_workflow +from ..services.proteindj_executor import launch_proteindj_workflow from ..services.proteinfold_executor import launch_proteinfold_workflow from ..services.seqera import WorkflowLaunchResult +from ..services.seqera_errors import SeqeraAPIError, SeqeraConfigurationError from ..services.wisps_executor import launch_wisps_workflow from . import SCHEDULER LAUNCH_MAX_ATTEMPTS = 3 RETRY_DELAY_BASE = 5 * 60 +# Gadi's gpuhopper PBS queue holds 50 job slots and each workflow run occupies approximately 2 of +# them (Nextflow's queueSize), so 25 workflows can run concurrently. Hardcoded as a +# temporary MVP value, configurable via env var. +MAX_CONCURRENT_WORKFLOWS = int(os.getenv("MAX_CONCURRENT_WORKFLOWS", "25")) + class LaunchFunction(Protocol): """ @@ -68,7 +76,12 @@ def launch_job(job_id: UUID, dry_run: bool = False) -> None: elif workflow_name in ("single-prediction", "proteinfold"): launch_func = launch_proteinfold_workflow elif workflow_name in ("de-novo-design", "bindflow", "bindcraft"): - launch_func = launch_bindflow_workflow + # de-novo-design covers two algorithms (bindcraft vs rfdiffusion), each + # with its own executor; workflow_run.tool holds the one selected at launch. + tool = (job.workflow_run.tool or "").lower() + launch_func = ( + launch_proteindj_workflow if tool == "rfdiffusion" else launch_bindflow_workflow + ) else: raise ValueError(f"Unsupported workflow: {job.workflow.name}") try: @@ -105,6 +118,20 @@ def launch_job(job_id: UUID, dry_run: bool = False) -> None: return +def get_available_workflow_capacity() -> int: + """ + How many more workflows can be submitted to Gadi right now, per the Seqera API's + count of workflows still occupying a job slot there (see MAX_CONCURRENT_WORKFLOWS). + """ + active_workflow_count = asyncio.run(seqera.count_active_workflows()) + capacity = max(0, MAX_CONCURRENT_WORKFLOWS - active_workflow_count) + logger.info( + f"{active_workflow_count}/{MAX_CONCURRENT_WORKFLOWS} workflows active on Gadi " + f"({capacity} submission slot(s) available)." + ) + return capacity + + def submit_pending_jobs(dry_run: bool = False): # Time between jobs job_offset = 10 @@ -115,6 +142,15 @@ def submit_pending_jobs(dry_run: bool = False): logger.warning("Skipping pending job submission while system status is unhealthy.") return + try: + available_capacity = get_available_workflow_capacity() + except (SeqeraAPIError, SeqeraConfigurationError) as e: + logger.warning(f"Could not determine Gadi workflow capacity from Seqera: {e}") + return + if available_capacity <= 0: + logger.info("Gadi is at its concurrent workflow limit; skipping submission this tick.") + return + now = datetime.now(tz=UTC) pending_query = select(QueuedJob).where( @@ -123,7 +159,13 @@ def submit_pending_jobs(dry_run: bool = False): pending_jobs = db_session.scalars(pending_query).all() logger.info(f"Found {len(pending_jobs)} pending jobs.") - for index, job in enumerate(pending_jobs): + jobs_to_submit = pending_jobs[:available_capacity] + if len(jobs_to_submit) < len(pending_jobs): + logger.info( + f"Only submitting {len(jobs_to_submit)} of {len(pending_jobs)} pending jobs " + "due to available Gadi capacity." + ) + for index, job in enumerate(jobs_to_submit): launch_id = f"launch_job_{job.id}" # Ignore if already scheduled if SCHEDULER.get_job(launch_id, jobstore="memory") is not None: diff --git a/app/schemas/health.py b/app/schemas/health.py index e95c67f..ee47da3 100644 --- a/app/schemas/health.py +++ b/app/schemas/health.py @@ -108,3 +108,41 @@ class SystemStatusAdminResponse(BaseModel): default=None, description="One-click link to the backend CloudWatch log group, if configured", ) + + +class ComponentIncident(BaseModel): + """A single downtime period for one component (never ``healthy``).""" + + status: HealthStatus + startedAt: datetime + endedAt: datetime | None = Field( + default=None, description="Null while the incident is still ongoing" + ) + message: str | None = None + + +class ComponentDowntimeSummary(BaseModel): + """Downtime aggregates for a component over the requested time window.""" + + incidentCount: int + downtimeSeconds: float + uptimePercent: float = Field( + description="Share of the window the component was healthy, as a percentage" + ) + + +class ComponentDowntime(BaseModel): + """Downtime incidents and summary for a single component.""" + + name: str + summary: ComponentDowntimeSummary + incidents: list[ComponentIncident] + + +class SystemStatusDowntimeResponse(BaseModel): + """Verbose per-component downtime returned by GET /admin/api/system-status/history.""" + + windowHours: int + since: datetime + until: datetime + components: list[ComponentDowntime] diff --git a/app/schemas/workflows.py b/app/schemas/workflows.py index 0dca971..f866266 100644 --- a/app/schemas/workflows.py +++ b/app/schemas/workflows.py @@ -2,16 +2,38 @@ from __future__ import annotations +import base64 +import re from datetime import datetime from enum import StrEnum -from typing import Any, Literal - -from pydantic import BaseModel, ConfigDict, Field, field_validator +from typing import Annotated, Any, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) WorkflowName = Literal[ "single-prediction", "de-novo-design", "bulk-prediction", "interaction-screening" ] WorkflowTool = Literal["alphafold2", "bindcraft", "boltz", "colabfold", "rfdiffusion"] +MoleculeType = Literal["protein", "rna", "dna", "ligand", "ccd"] + +NonEmptyStr = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] + +SINGLE_PREDICTION_MAX_ENTITIES = 52 +SINGLE_PREDICTION_LIGAND_SIZE = 30 +SINGLE_PREDICTION_SIZE_LIMITS: dict[str, int] = { + "alphafold2": 2000, + "colabfold": 4000, + "boltz": 4000, +} +SINGLE_PREDICTION_BOLTZ_POTENTIALS_SIZE_LIMIT = 2000 class PipelineStatus(StrEnum): @@ -97,10 +119,20 @@ class WorkflowUserDetails(BaseModel): """ user_email: str = Field(..., description="Email address of the user") - full_name: str = Field(..., description="Full name of the user") - institute: str = Field(..., description="Institute of the user") ip_address: str = Field(..., description="IP address of the user") + def get_encoded_account_details(self) -> str: + """Return the `-A` account string for cluster job submission. + + Encodes email and IP as base64 (colon-joined) so the account string + never leaks the user's raw email/IP into cluster logs or job metadata. + """ + encoded_email = base64.b64encode(self.user_email.encode()).decode() + if not self.ip_address: + return encoded_email + encoded_ip = base64.b64encode(self.ip_address.encode()).decode() + return f"{encoded_email}:{encoded_ip}" + class WispsFormData(WorkflowFormData): """Form data for WISPS workflows (interaction-screening, bulk-prediction).""" @@ -113,6 +145,117 @@ class WispsFormData(WorkflowFormData): ) +class SinglePredictionEntity(BaseModel): + """A single entity row submitted for a single-prediction run.""" + + model_config = ConfigDict(extra="allow") + + id: NonEmptyStr | None = None + moleculeType: MoleculeType + copyNumber: PositiveInt = 1 + sequence: Annotated[str, StringConstraints(strip_whitespace=True)] = "" + + +def _entity_prediction_size(entity: SinglePredictionEntity) -> int: + """Prediction size for one copy of an entity (ligands use a fixed size).""" + if entity.moleculeType in ("ligand", "ccd"): + return SINGLE_PREDICTION_LIGAND_SIZE + return len(re.sub(r"\s+", "", entity.sequence or "")) + + +def single_prediction_size_limit(tool: str, boltz_use_potentials: bool) -> int: + """Exclusive upper bound on the total prediction size for the given tool.""" + if tool == "boltz" and boltz_use_potentials: + return SINGLE_PREDICTION_BOLTZ_POTENTIALS_SIZE_LIMIT + return SINGLE_PREDICTION_SIZE_LIMITS.get(tool, SINGLE_PREDICTION_SIZE_LIMITS["colabfold"]) + + +def validate_single_prediction_entities( + entities: list[SinglePredictionEntity], + tool: str, + *, + boltz_use_potentials: bool = False, +) -> None: + """Validate entity count, protein presence, and total prediction size. + + Raises ``ValueError`` with a user-facing message on the first violation. + """ + if not entities: + raise ValueError("At least one entity is required.") + + total_copies = 0 + total_size = 0 + has_protein = False + for entity in entities: + total_copies += entity.copyNumber + if entity.moleculeType == "protein": + has_protein = True + total_size += _entity_prediction_size(entity) * entity.copyNumber + + if total_copies > SINGLE_PREDICTION_MAX_ENTITIES: + raise ValueError( + f"Too many entities: {total_copies} including copies. " + f"The maximum allowed is {SINGLE_PREDICTION_MAX_ENTITIES}." + ) + + if not has_protein: + raise ValueError( + "At least one entity must be a protein. " + "This workflow cannot run without a protein input." + ) + + limit = single_prediction_size_limit(tool, boltz_use_potentials) + if total_size >= limit: + raise ValueError( + f"Total prediction size ({total_size}) must be less than {limit} for {tool}." + ) + + +MAX_HOTSPOT_RESIDUES = 8 +MIN_DESIGN_LENGTH = 65 +MAX_DESIGN_LENGTH = 150 + + +class ProteinDjFormData(WorkflowFormData): + """Form data for the ProteinDJ (rfdiffusion) de-novo-design workflow.""" + + starting_pdb: str = Field(..., description="S3 URI of the uploaded starting PDB file") + target_hotspot_residues: str = Field( + ..., description="Comma-separated hotspot residues, e.g. 'A20,A21'" + ) + number_of_final_designs: int = Field(..., ge=1, description="Number of designs to generate") + min_length: int = Field(..., ge=MIN_DESIGN_LENGTH, description="Minimum binder length") + max_length: int = Field(..., le=MAX_DESIGN_LENGTH, description="Maximum binder length") + + @field_validator("starting_pdb") + @classmethod + def _validate_starting_pdb(cls, value: str) -> str: + stripped = value.strip() + if not stripped: + raise ValueError("This field is required") + return stripped + + @field_validator("target_hotspot_residues") + @classmethod + def _validate_hotspot_residues(cls, value: str) -> str: + stripped = value.strip() + if not stripped: + raise ValueError("This field is required") + residue_count = len([token for token in stripped.split(",") if token.strip()]) + if residue_count > MAX_HOTSPOT_RESIDUES: + raise ValueError( + f"Too many hotspot residues ({residue_count}); only up to " + f"{MAX_HOTSPOT_RESIDUES} are supported" + ) + return stripped + + @model_validator(mode="after") + def _validate_length_range(self) -> ProteinDjFormData: + if self.min_length > self.max_length: + raise ValueError("min_length must not exceed max_length") + return self + + class WorkflowLaunchPayload(BaseModel): model_config = ConfigDict(extra="forbid") diff --git a/app/services/bindflow_config.py b/app/services/bindflow_config.py index 750bde2..8f831d5 100644 --- a/app/services/bindflow_config.py +++ b/app/services/bindflow_config.py @@ -5,13 +5,14 @@ from typing import Any from ..schemas.workflows import WorkflowUserDetails +from .cluster_utils import GADI_PROJECT from .workflow_config_fetcher import fetch_workflow_config def get_bindflow_default_params(out_dir: str, samplesheet_url: str) -> dict[str, Any]: """Get default parameters for bindflow workflow.""" return { - "project": "yz52", + "project": GADI_PROJECT, "outdir": out_dir, "input": samplesheet_url, } @@ -25,17 +26,11 @@ def get_bindflow_config_profiles() -> list[str]: def get_bindflow_config_text( config_file_path: str, *, - job_id: str, user_details: WorkflowUserDetails, - timestamp: str, ) -> str: """Read bindflow base config and append a process override block with runtime values.""" base = fetch_workflow_config(config_file_path) - cluster_opts = ( - f"-v JOB_ID={job_id},USER_NAME={user_details.user_email}," - f"TIMESTAMP={timestamp},FULL_NAME={user_details.full_name}," - f"INSTITUTE={user_details.institute},IP_ADDRESS={user_details.ip_address}" - ) + cluster_opts = f"-A {user_details.get_encoded_account_details()}" override = f'\nprocess {{\n clusterOptions = "{cluster_opts}"\n}}\n' return base + override diff --git a/app/services/bindflow_executor.py b/app/services/bindflow_executor.py index 9296160..7f75a92 100644 --- a/app/services/bindflow_executor.py +++ b/app/services/bindflow_executor.py @@ -16,7 +16,12 @@ get_bindflow_config_text, get_bindflow_default_params, ) -from .launch_payloads import get_executor_script, inject_prerun_script, without_prerun_script +from .launch_payloads import ( + DEFAULT_MODULE_LOADS, + get_executor_script, + inject_prerun_script, + without_prerun_script, +) from .seqera import ( WorkflowLaunchResult, _get_required_env, @@ -58,14 +63,9 @@ async def prepare_bindflow_workflow( # pylint: disable=too-many-locals raise SeqeraConfigurationError("Missing output identifier for workflow launch") out_dir = f"s3://{s3_bucket}/{output_key}" - timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") - dataset_url = f"s3://{s3_bucket}/{s3_input_key}" default_params = get_bindflow_default_params(out_dir, dataset_url) - default_params["job_id"] = run_name - default_params["user_name"] = user_details.user_email - default_params["timestamp"] = timestamp default_params["mode"] = mode # Merge any tool-specific params forwarded from the frontend form @@ -91,9 +91,7 @@ async def prepare_bindflow_workflow( # pylint: disable=too-many-locals "configProfiles": get_bindflow_config_profiles(), "configText": get_bindflow_config_text( config_path, - job_id=run_name, user_details=user_details, - timestamp=timestamp, ), "resume": False, } @@ -136,7 +134,7 @@ async def launch_bindflow_workflow( # pylint: disable=too-many-locals prerun_script = get_executor_script( prerun_script_path=queued_job.workflow.prerun_script_path, - module_loads=["singularity", "nextflow"], + module_loads=DEFAULT_MODULE_LOADS, env={ "AWS_ACCESS_KEY_ID": os.getenv("AWS_ACCESS_KEY_ID", ""), "AWS_SECRET_ACCESS_KEY": os.getenv("AWS_SECRET_ACCESS_KEY", ""), diff --git a/app/services/cluster_utils.py b/app/services/cluster_utils.py new file mode 100644 index 0000000..b7d3da2 --- /dev/null +++ b/app/services/cluster_utils.py @@ -0,0 +1,7 @@ +"""Utilities for building HPC cluster option strings.""" + +from __future__ import annotations + +import os + +GADI_PROJECT: str = os.getenv("GADI_PROJECT", "yz52") diff --git a/app/services/health.py b/app/services/health.py index a8785a4..0e58123 100644 --- a/app/services/health.py +++ b/app/services/health.py @@ -23,10 +23,10 @@ from urllib.parse import quote import httpx -from sqlalchemy import func, select +from sqlalchemy import func, or_, select from sqlalchemy.orm import Session -from ..db.models.system_status import SystemStatusCache +from ..db.models.system_status import SystemStatusCache, SystemStatusIncident from ..schemas.health import ( COMPONENT_COMPUTE_ENV, COMPONENT_SEQERA_API, @@ -487,8 +487,57 @@ def _try_acquire_cache_refresh_lock(db: Session) -> bool: return bool(db.execute(select(func.pg_try_advisory_xact_lock(_CACHE_LOCK_ID))).scalar_one()) +def _open_incident(db: Session, status: SystemStatus, component: ProbeResult) -> None: + db.add( + SystemStatusIncident( + component=component.name, + status=component.status, + started_at=status.checked_at, + message=component.message, + ) + ) + + +def _update_incidents(db: Session, status: SystemStatus) -> None: + """Open/close/split incident rows so only downtime periods are stored. + + Healthy checks are never recorded. A component's most recent incident with + ``ended_at IS NULL`` (if any) is its currently-open outage: + + - Recovers to healthy -> close it. + - Still down at the same severity -> leave it open (refresh the message). + - Still down but severity changed (e.g. degraded -> unhealthy) -> close it + and open a new one, so each row represents a single severity level. + - No open incident and still down -> open a new one. + """ + for component in status.components: + open_incident = db.execute( + select(SystemStatusIncident) + .where(SystemStatusIncident.component == component.name) + .where(SystemStatusIncident.ended_at.is_(None)) + .order_by(SystemStatusIncident.started_at.desc()) + .limit(1) + ).scalar_one_or_none() + + if component.status == "healthy": + if open_incident is not None: + open_incident.ended_at = status.checked_at + db.add(open_incident) + continue + + if open_incident is None: + _open_incident(db, status, component) + elif open_incident.status != component.status: + open_incident.ended_at = status.checked_at + db.add(open_incident) + _open_incident(db, status, component) + elif component.message != open_incident.message: + open_incident.message = component.message + db.add(open_incident) + + async def refresh_db_cache(db: Session) -> SystemStatus: - """Run probes and replace the single shared DB cache row.""" + """Run probes, replace the single shared DB cache row, and log incidents.""" status = await collect_system_status() now = datetime.now(UTC) expires_at = now + timedelta(seconds=_CACHE_TTL_SECONDS) @@ -510,6 +559,7 @@ async def refresh_db_cache(db: Session) -> SystemStatus: row.updated_at = now db.add(row) + _update_incidents(db, status) db.commit() return status @@ -610,6 +660,95 @@ def to_public_dict(status: SystemStatus) -> dict[str, Any]: } +def _as_utc(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=UTC) + + +def get_incidents( + db: Session, + *, + since: datetime, + until: datetime, + component: str | None = None, +) -> list[SystemStatusIncident]: + """Return incidents overlapping ``[since, until]``, oldest first. + + An incident overlaps the window if it started at/before ``until`` and + either is still open (``ended_at IS NULL``) or ended at/after ``since``. + """ + stmt = select(SystemStatusIncident).where( + SystemStatusIncident.started_at <= until, + or_( + SystemStatusIncident.ended_at.is_(None), + SystemStatusIncident.ended_at >= since, + ), + ) + if component: + stmt = stmt.where(SystemStatusIncident.component == component) + stmt = stmt.order_by(SystemStatusIncident.component, SystemStatusIncident.started_at) + return list(db.execute(stmt).scalars()) + + +def _clipped_downtime_seconds( + incident: SystemStatusIncident, *, since: datetime, until: datetime +) -> float: + """Downtime contributed by ``incident`` within ``[since, until]``. + + Incidents that started before ``since`` or are still ongoing (``ended_at`` + is None, so it counts as down through ``until``) are clipped to the window + so a long-running outage doesn't over-count time outside the query range. + """ + start = max(_as_utc(incident.started_at), since) + end = min(_as_utc(incident.ended_at) if incident.ended_at is not None else until, until) + return max((end - start).total_seconds(), 0.0) + + +def _summarize_downtime( + incidents: list[SystemStatusIncident], *, since: datetime, until: datetime +) -> dict[str, Any]: + downtime_seconds = sum( + _clipped_downtime_seconds(incident, since=since, until=until) for incident in incidents + ) + window_seconds = max((until - since).total_seconds(), 0.0) + uptime_percent = ( + round(max(window_seconds - downtime_seconds, 0.0) / window_seconds * 100, 2) + if window_seconds + else 100.0 + ) + return { + "incidentCount": len(incidents), + "downtimeSeconds": round(downtime_seconds, 2), + "uptimePercent": uptime_percent, + } + + +def to_downtime_dict( + incidents: list[SystemStatusIncident], *, since: datetime, until: datetime, hours: int +) -> dict[str, Any]: + """Group flat incident rows by component into the admin downtime response shape.""" + by_component: dict[str, list[SystemStatusIncident]] = {} + for incident in incidents: + by_component.setdefault(incident.component, []).append(incident) + + components = [ + { + "name": name, + "summary": _summarize_downtime(rows, since=since, until=until), + "incidents": [ + { + "status": row.status, + "startedAt": row.started_at, + "endedAt": row.ended_at, + "message": row.message, + } + for row in rows + ], + } + for name, rows in by_component.items() + ] + return {"windowHours": hours, "since": since, "until": until, "components": components} + + def to_admin_dict(status: SystemStatus) -> dict[str, Any]: """Verbose projection: includes raw probe detail and a CloudWatch link.""" return { diff --git a/app/services/job_utils.py b/app/services/job_utils.py index 3015d95..c01e775 100644 --- a/app/services/job_utils.py +++ b/app/services/job_utils.py @@ -206,3 +206,27 @@ async def ensure_completed_run_score(db: Session, run: WorkflowRun, ui_status: s db.add(RunMetric(run_id=run.id, max_score=bounded_score)) db.commit() return _round_score(bounded_score) + + +async def sync_service_usage(db: Session, run: WorkflowRun, ui_status: str) -> float | None: + if ui_status != "Completed": + return None + if run.service_usage is not None: + return run.service_usage + + try: + spec = get_output_spec(run) + except ValueError as exc: + logger.warning(f"Can't get service usage for run {run.id} - can't get output spec: {exc}") + return None + + await sync_workflow_outputs(db, run=run, spec=spec, suppress_s3_errors=True) + + usage = await spec.get_service_units(db, run) + if usage is None: + return None + else: + run.service_usage = usage + db.add(run) + db.commit() + return usage diff --git a/app/services/launch_payloads.py b/app/services/launch_payloads.py index 130ff7a..6bdeb79 100644 --- a/app/services/launch_payloads.py +++ b/app/services/launch_payloads.py @@ -7,6 +7,8 @@ from .workflow_config_fetcher import fetch_workflow_config +DEFAULT_MODULE_LOADS = ["singularity", "nextflow/25.10.4"] + def without_prerun_script(launch_payload: dict[str, Any]) -> dict[str, Any]: """Return a copy safe to persist in the job queue.""" diff --git a/app/services/proteindj_config.py b/app/services/proteindj_config.py new file mode 100644 index 0000000..732711c --- /dev/null +++ b/app/services/proteindj_config.py @@ -0,0 +1,47 @@ +"""ProteinDJ workflow configuration and executor settings (modeled after bindflow).""" + +from __future__ import annotations + +from typing import Any + +from ..schemas.workflows import WorkflowUserDetails +from .workflow_config_fetcher import fetch_workflow_config + + +def get_proteindj_default_params( + out_dir: str, + input_pdb: str, + hotspot_residues: str, + num_designs: int, + design_length: str, +) -> dict[str, Any]: + """Get default parameters for proteindj workflow. + + ProteinDJ (rfdiffusion) takes a single PDB plus design params directly — + no samplesheet — so these are passed straight through as paramsText keys. + """ + return { + "out_dir": out_dir, + "input_pdb": input_pdb, + "hotspot_residues": hotspot_residues, + "num_designs": num_designs, + "design_length": design_length, + } + + +def get_proteindj_config_profiles() -> list[str]: + """Get config profiles for proteindj workflow.""" + return ["singularity"] + + +def get_proteindj_config_text( + config_file_path: str, + *, + user_details: WorkflowUserDetails, +) -> str: + """Read proteindj base config and append a process override block with runtime values.""" + base = fetch_workflow_config(config_file_path) + + cluster_opts = f"-A {user_details.get_encoded_account_details()}" + override = f'\nprocess {{\n clusterOptions = "{cluster_opts}"\n}}\n' + return base + override diff --git a/app/services/proteindj_executor.py b/app/services/proteindj_executor.py new file mode 100644 index 0000000..345d073 --- /dev/null +++ b/app/services/proteindj_executor.py @@ -0,0 +1,177 @@ +"""ProteinDJ workflow executor for Seqera Platform (modeled after bindflow).""" + +from __future__ import annotations + +import logging +import os +from datetime import UTC, datetime +from typing import Any + +from pydantic import ValidationError +from sqlalchemy.orm import Session + +from ..db.models import QueuedJob, WorkflowRun +from ..schemas.workflows import ( + ProteinDjFormData, + WorkflowFormData, + WorkflowLaunchForm, + WorkflowUserDetails, +) +from .launch_payloads import ( + DEFAULT_MODULE_LOADS, + get_executor_script, + inject_prerun_script, + without_prerun_script, +) +from .proteindj_config import ( + get_proteindj_config_profiles, + get_proteindj_config_text, + get_proteindj_default_params, +) +from .seqera import ( + WorkflowLaunchResult, + _get_required_env, + params_to_yaml_text, + post_seqera_launch, +) +from .seqera_errors import SeqeraConfigurationError + +logger = logging.getLogger(__name__) + + +def _design_length(fields: ProteinDjFormData) -> str: + # The frontend's Input Configuration step (shared with bindcraft) sends + # min_length/max_length as separate fields; ProteinDJ expects them as a + # single "min-max" range. + return f"{fields.min_length}-{fields.max_length}" + + +def _parse_proteindj_form_data(form_data: WorkflowFormData) -> ProteinDjFormData: + try: + return ProteinDjFormData.model_validate(form_data.model_dump()) + except ValidationError as exc: + missing = "formData" + for error in exc.errors(): + loc = error.get("loc") + if loc: + *_, field_name = loc # last element of the location path is the field name + missing = str(field_name) + break + raise SeqeraConfigurationError( + f"'{missing}' is required in formData for ProteinDJ workflow launch" + ) from exc + + +async def prepare_proteindj_workflow( # pylint: disable=too-many-locals + form: WorkflowLaunchForm, + *, + db_session: Session, + workflow_run: WorkflowRun, + pipeline: str, + config_path: str, + revision: str | None = None, + output_id: str | None = None, + form_data: WorkflowFormData, + user_details: WorkflowUserDetails, + commit: bool = False, +) -> QueuedJob: + """Build and queue a proteindj launch payload.""" + workspace_id = _get_required_env("WORK_SPACE") + compute_env_id = _get_required_env("COMPUTE_ID") + work_dir = _get_required_env("WORK_DIR") + s3_bucket = _get_required_env("AWS_S3_BUCKET") + + run_name = (form.runName or "").strip() + if not run_name: + raise SeqeraConfigurationError("Missing run name for workflow launch") + # Always use a unique backend-generated ID for outputs to avoid S3 prefix collisions. + output_key = (output_id or "").strip() + if not output_key: + raise SeqeraConfigurationError("Missing output identifier for workflow launch") + out_dir = f"s3://{s3_bucket}/{output_key}" + + proteindj_fields = _parse_proteindj_form_data(form_data) + default_params = get_proteindj_default_params( + out_dir, + input_pdb=proteindj_fields.starting_pdb, + hotspot_residues=proteindj_fields.target_hotspot_residues, + num_designs=proteindj_fields.number_of_final_designs, + design_length=_design_length(proteindj_fields), + ) + + # Serialize to YAML + params_text = params_to_yaml_text(default_params) + + # Add custom paramsText from frontend if provided + if form.paramsText and form.paramsText.strip(): + params_text = f"{params_text}\n{form.paramsText.rstrip()}" + + launch_payload: dict[str, Any] = { + "computeEnvId": compute_env_id, + "runName": run_name, + "pipeline": pipeline, + "workDir": work_dir, + "workspaceId": workspace_id, + "revision": revision or "dev", + "paramsText": params_text, + "configProfiles": get_proteindj_config_profiles(), + "configText": get_proteindj_config_text( + config_path, + user_details=user_details, + ), + "resume": False, + } + + queued_job = QueuedJob( + workflow=workflow_run.workflow, + workflow_run=workflow_run, + launch_payload=without_prerun_script(launch_payload), + status="pending", + next_attempt_at=datetime.now(UTC), + ) + db_session.add(queued_job) + if commit: + db_session.commit() + else: + db_session.flush() + return queued_job + + +async def launch_proteindj_workflow( # pylint: disable=too-many-locals + *, + queued_job: QueuedJob, + dry_run: bool = False, +) -> WorkflowLaunchResult | None: + """Launch a proteindj workflow on the Seqera Platform.""" + launch_payload = queued_job.launch_payload + + # Log the complete params being sent + logger.info("Launch payload paramsText", extra={"paramsText": launch_payload["paramsText"]}) + + logger.info( + "Launching proteindj workflow via Seqera API", + extra={ + "workspaceId": launch_payload["workspaceId"], + "computeEnvId": launch_payload["computeEnvId"], + "pipeline": launch_payload["pipeline"], + "runName": launch_payload["runName"], + }, + ) + + prerun_script = get_executor_script( + prerun_script_path=queued_job.workflow.prerun_script_path, + module_loads=DEFAULT_MODULE_LOADS, + env={ + "AWS_ACCESS_KEY_ID": os.getenv("AWS_ACCESS_KEY_ID", ""), + "AWS_SECRET_ACCESS_KEY": os.getenv("AWS_SECRET_ACCESS_KEY", ""), + "AWS_REGION": os.getenv("AWS_REGION", "ap-southeast-2"), + }, + ) + runtime_payload = inject_prerun_script( + launch_payload=launch_payload, prerun_script=prerun_script + ) + + if dry_run: + logger.info("Dry run - not launching proteindj workflow") + return None + return await post_seqera_launch({"launch": runtime_payload}, workflow_label="ProteinDJ") diff --git a/app/services/proteinfold_config.py b/app/services/proteinfold_config.py index 8843a13..b836a2a 100644 --- a/app/services/proteinfold_config.py +++ b/app/services/proteinfold_config.py @@ -5,6 +5,7 @@ from typing import Any from ..schemas.workflows import WorkflowUserDetails +from .cluster_utils import GADI_PROJECT from .workflow_config_fetcher import fetch_workflow_config @@ -12,7 +13,7 @@ def get_proteinfold_default_params( out_dir: str, samplesheet_url: str, mode: str = "alphafold2" ) -> dict[str, Any]: """Get default parameters for proteinfold workflow.""" - return {"input": samplesheet_url, "outdir": out_dir, "project": "yz52", "mode": mode} + return {"input": samplesheet_url, "outdir": out_dir, "project": GADI_PROJECT, "mode": mode} def get_proteinfold_config_profiles() -> list[str]: @@ -23,17 +24,11 @@ def get_proteinfold_config_profiles() -> list[str]: def get_proteinfold_config_text( config_file_path: str, *, - job_id: str, user_details: WorkflowUserDetails, - timestamp: str, ) -> str: """Read proteinfold base config and append a process override block with runtime values.""" base = fetch_workflow_config(config_file_path) - cluster_opts = ( - f"-P yz52 -v JOB_ID={job_id},USER_NAME={user_details.user_email}," - f"TIMESTAMP={timestamp},FULL_NAME={user_details.full_name}," - f"INSTITUTE={user_details.institute},IP_ADDRESS={user_details.ip_address}" - ) + cluster_opts = f"-P {GADI_PROJECT} -A {user_details.get_encoded_account_details()}" override = f'\nprocess {{\n clusterOptions = "{cluster_opts}"\n}}\n' return base + override diff --git a/app/services/proteinfold_executor.py b/app/services/proteinfold_executor.py index cfb99cc..55886f0 100644 --- a/app/services/proteinfold_executor.py +++ b/app/services/proteinfold_executor.py @@ -11,7 +11,12 @@ from ..db.models import QueuedJob, WorkflowRun from ..schemas.workflows import WorkflowFormData, WorkflowLaunchForm, WorkflowUserDetails -from .launch_payloads import get_executor_script, inject_prerun_script, without_prerun_script +from .launch_payloads import ( + DEFAULT_MODULE_LOADS, + get_executor_script, + inject_prerun_script, + without_prerun_script, +) from .proteinfold_config import ( get_proteinfold_config_profiles, get_proteinfold_config_text, @@ -30,7 +35,7 @@ # Params forwarded from the frontend's Tool Settings (step 2) _TOOL_PARAM_KEYS = frozenset( { - "alphafold2_random_seed", + "random_seed", "alphafold2_full_dbs", "colabfold_num_recycles", "colabfold_use_templates", @@ -50,14 +55,11 @@ def _build_params_text( mode: str, form_data: WorkflowFormData | None, custom_params: str | None, - extra_params: dict[str, Any] | None = None, ) -> str: """Build the YAML params string for the Seqera launch payload.""" params = get_proteinfold_default_params(out_dir, samplesheet_url, mode) if form_data: params.update(_tool_params(form_data)) - if extra_params: - params.update(extra_params) params_text = params_to_yaml_text(params) if custom_params and custom_params.strip(): params_text = f"{params_text}\n{custom_params.rstrip()}" @@ -89,9 +91,7 @@ async def prepare_proteinfold_workflow( raise SeqeraConfigurationError("Missing output identifier for workflow launch") out_dir = f"s3://{s3_bucket}/{output_id.strip()}" - timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") - job_id = (form.runName or "").strip() - if not job_id: + if not form.runName or not form.runName.strip(): raise SeqeraConfigurationError("Missing run name for workflow launch") sheet_url = f"s3://{s3_bucket}/{s3_input_key}" @@ -101,11 +101,6 @@ async def prepare_proteinfold_workflow( mode, form_data, form.paramsText, - extra_params={ - "job_id": job_id, - "user_name": user_details.user_email, - "timestamp": timestamp, - }, ) launch_payload: dict[str, Any] = { @@ -119,9 +114,7 @@ async def prepare_proteinfold_workflow( "configProfiles": get_proteinfold_config_profiles(), "configText": get_proteinfold_config_text( config_path, - job_id=job_id, user_details=user_details, - timestamp=timestamp, ), "resume": False, } @@ -161,7 +154,7 @@ async def launch_proteinfold_workflow( prerun_script = get_executor_script( prerun_script_path=queued_job.workflow.prerun_script_path, - module_loads=["singularity", "nextflow"], + module_loads=DEFAULT_MODULE_LOADS, env={ "AWS_ACCESS_KEY_ID": os.getenv("AWS_ACCESS_KEY_ID", ""), "AWS_SECRET_ACCESS_KEY": os.getenv("AWS_SECRET_ACCESS_KEY", ""), diff --git a/app/services/results_utils.py b/app/services/results_utils.py index 9276a81..0f98c7b 100644 --- a/app/services/results_utils.py +++ b/app/services/results_utils.py @@ -27,7 +27,7 @@ read_s3_file, ) -OutputCategory = Literal["report", "stats_csv", "pdb", "snapshot", "alignment"] +OutputCategory = Literal["report", "stats_csv", "pdb", "snapshot", "alignment", "usage"] class OutputClassifier(Protocol): @@ -59,11 +59,22 @@ class WorkflowResultsSpec: tool: WorkflowTool required_categories: set[OutputCategory] get_prefixes: Callable[[WorkflowRun], list[str]] - classify: OutputClassifier + classifier: OutputClassifier get_score_file: GetScoreFile extract_max_score: Callable[[str], Awaitable[float | None]] supports_snapshots: bool = False + def classify_output(self, key: str, sample_id: str | None) -> ClassifiedOutput | None: + """ + Classify a workflow output into a category. + First checks for outputs that are shared across workflows, + then the unique outputs for the specific workflow. + """ + shared_output = classify_shared_outputs(key) + if shared_output is not None: + return shared_output + return self.classifier(key, sample_id) + async def get_max_score(self, db: Session, run: WorkflowRun): keys = _get_run_output_keys(db, run) sample_id = get_sample_id_for_result(run) @@ -76,6 +87,27 @@ async def get_max_score(self, db: Session, run: WorkflowRun): logger.warning("Failed to extract max score from %r: %s", score_file, e, exc_info=True) return None + def get_usage_file(self, keys: list[str]) -> str | None: + for key in keys: + normalized = key.strip() + if not normalized: + continue + basename = normalized.rsplit("/", 1)[-1] + if basename == "UsageReport.csv": + return normalized + return None + + async def get_service_units(self, db: Session, run: WorkflowRun) -> float | None: + keys = _get_run_output_keys(db, run) + usage_file = self.get_usage_file(keys) + if usage_file is None: + return None + try: + return await get_run_service_usage(usage_file) + except Exception as e: + logger.warning("Failed to get service usage from %r: %s", usage_file, e, exc_info=True) + return None + _LOG_LEVEL_PATTERN = re.compile(r"\b(TRACE|DEBUG|INFO|WARN|WARNING|ERROR|FATAL)\b") _LOG_TIMESTAMP_PATTERN = re.compile( @@ -88,6 +120,22 @@ async def get_max_score(self, db: Session, run: WorkflowRun): logger = logging.getLogger(__name__) +def classify_shared_outputs(key: str) -> ClassifiedOutput | None: + """ + Classify outputs that are shared across workflows. Currently + only the usage report, but other common outputs can be added + """ + normalized = key.strip() + if not normalized or normalized.endswith("/"): + return None + + basename = normalized.rsplit("/", 1)[-1] + if basename == "UsageReport.csv": + return ClassifiedOutput(category="usage", label=basename) + + return None + + def _sanitize_content_disposition_filename(filename: str) -> str: sanitized = _HEADER_UNSAFE_FILENAME_CHARS.sub("_", filename).strip() return sanitized or "download" @@ -474,7 +522,7 @@ def classify_boltz_proteinfold_output( sample_id_pattern = re.escape(sample_id) if sample_id else "single_prediction" return classify_proteinfold_output_key( key, - pdb_pattern=rf"/boltz/top_ranked_structures/{sample_id_pattern}\.pdb", + pdb_pattern=rf"/boltz/top_ranked_structures/{sample_id_pattern}\.(?:cif|pdb)", # Find across all subfolders stats_pattern=rf"/boltz/{sample_id_pattern}/.+\.tsv", alignment_pattern=rf"/mmseqs/{sample_id_pattern}\.a3m", @@ -511,6 +559,7 @@ def build_bindcraft_output_listing_prefixes(run: WorkflowRun) -> list[str]: # Always include run-UUID-only prefixes; these do not depend on sample_id. prefixes: list[str] = [ + f"{run_uuid}/", f"{run_uuid}/ranker/", f"{run_uuid}/generate/", ] @@ -535,6 +584,7 @@ def build_boltz_proteinfold_output_listing_prefixes(run: WorkflowRun) -> list[st sample_id = get_sample_id_for_result(run) prefixes = [ + f"{run_uuid}/", f"{run_uuid}/reports/", f"{run_uuid}/boltz/top_ranked_structures/", f"{run_uuid}/mmseqs/", @@ -556,6 +606,7 @@ def build_alphafold2_proteinfold_output_listing_prefixes(run: WorkflowRun) -> li sample_id = get_sample_id_for_result(run) prefixes = [ + f"{run_uuid}/", f"{run_uuid}/reports/", f"{run_uuid}/alphafold2/split_msa_prediction/top_ranked_structures/", ] @@ -627,6 +678,7 @@ def build_wisps_output_listing_prefixes(run: WorkflowRun) -> list[str]: if not run_uuid: return [] return [ + f"{run_uuid}/", f"{run_uuid}/multiqc/", f"{run_uuid}/collect/", f"{run_uuid}/ipsae/", @@ -641,6 +693,7 @@ def build_colabfold_proteinfold_output_listing_prefixes(run: WorkflowRun) -> lis sample_id = get_sample_id_for_result(run) prefixes = [ + f"{run_uuid}/", f"{run_uuid}/reports/", f"{run_uuid}/colabfold/top_ranked_structures/", f"{run_uuid}/mmseqs/", @@ -654,6 +707,52 @@ def build_colabfold_proteinfold_output_listing_prefixes(run: WorkflowRun) -> lis return prefixes +def classify_rfdiffusion_output_key( + key: str, sample_id: str | None = None +) -> ClassifiedOutput | None: + normalized = key.strip() + if not normalized or normalized.endswith("/"): + return None + basename = normalized.rsplit("/", 1)[-1] + if "/results/" not in normalized.lower(): + return None + if basename == "ranked_designs.csv": + return ClassifiedOutput(category="stats_csv", label=basename) + if basename == "ranked_designs.tar.gz": + return ClassifiedOutput(category="pdb", label=basename) + return None + + +def get_rfdiffusion_score_file(keys: list[str], sample_id: str | None) -> str | None: + for key in keys: + normalized = key.strip() + if not normalized: + continue + basename = normalized.rsplit("/", 1)[-1] + if basename == "ranked_designs.csv": + return normalized + return None + + +async def extract_rfdiffusion_max_score(score_file: str) -> float | None: + content = await read_s3_file(score_file) + csv_reader = csv.DictReader(StringIO(content)) + row = next(csv_reader, None) + if row is None: + return None + value = row.get("af2_plddt_overall") + if value and value.strip(): + return float(value) / 100 + return None + + +def build_rfdiffusion_output_listing_prefixes(run: WorkflowRun) -> list[str]: + run_uuid = str(getattr(run, "id", "") or "").strip() + if not run_uuid: + return [] + return [f"{run_uuid}/", f"{run_uuid}/results/"] + + def _make_wisps_spec(tool: WorkflowTool) -> WorkflowResultsSpec: return WorkflowResultsSpec( kind="interaction-screening", @@ -662,7 +761,7 @@ def _make_wisps_spec(tool: WorkflowTool) -> WorkflowResultsSpec: get_prefixes=build_wisps_output_listing_prefixes, get_score_file=get_wisps_score_file, extract_max_score=extract_wisps_max_score, - classify=classify_wisps_output_key, + classifier=classify_wisps_output_key, ) @@ -674,7 +773,7 @@ def _make_bulk_prediction_spec(tool: WorkflowTool) -> WorkflowResultsSpec: get_prefixes=build_wisps_output_listing_prefixes, get_score_file=get_wisps_score_file, extract_max_score=extract_bulk_prediction_max_score, - classify=classify_wisps_output_key, + classifier=classify_wisps_output_key, ) @@ -687,9 +786,18 @@ def _make_bulk_prediction_spec(tool: WorkflowTool) -> WorkflowResultsSpec: get_prefixes=build_bindcraft_output_listing_prefixes, get_score_file=get_bindcraft_score_file, extract_max_score=extract_bindcraft_max_score, - classify=classify_bindcraft_output_key, + classifier=classify_bindcraft_output_key, supports_snapshots=True, ), + "rfdiffusion": WorkflowResultsSpec( + kind="de-novo-design", + tool="rfdiffusion", + required_categories={"stats_csv", "pdb"}, + get_prefixes=build_rfdiffusion_output_listing_prefixes, + get_score_file=get_rfdiffusion_score_file, + extract_max_score=extract_rfdiffusion_max_score, + classifier=classify_rfdiffusion_output_key, + ), }, "single-prediction": { "boltz": WorkflowResultsSpec( @@ -699,7 +807,7 @@ def _make_bulk_prediction_spec(tool: WorkflowTool) -> WorkflowResultsSpec: get_prefixes=build_boltz_proteinfold_output_listing_prefixes, get_score_file=get_proteinfold_score_file, extract_max_score=extract_proteinfold_max_score, - classify=classify_boltz_proteinfold_output, + classifier=classify_boltz_proteinfold_output, ), "alphafold2": WorkflowResultsSpec( kind="single-prediction", @@ -708,7 +816,7 @@ def _make_bulk_prediction_spec(tool: WorkflowTool) -> WorkflowResultsSpec: get_prefixes=build_alphafold2_proteinfold_output_listing_prefixes, get_score_file=get_proteinfold_score_file, extract_max_score=extract_proteinfold_max_score, - classify=classify_alphafold2_proteinfold_output, + classifier=classify_alphafold2_proteinfold_output, ), "colabfold": WorkflowResultsSpec( kind="single-prediction", @@ -717,7 +825,7 @@ def _make_bulk_prediction_spec(tool: WorkflowTool) -> WorkflowResultsSpec: get_prefixes=build_colabfold_proteinfold_output_listing_prefixes, get_score_file=get_proteinfold_score_file, extract_max_score=extract_proteinfold_max_score, - classify=classify_colabfold_proteinfold_output, + classifier=classify_colabfold_proteinfold_output, ), }, "interaction-screening": { @@ -747,7 +855,7 @@ def collect_classified_outputs( outputs = {} sample_id = get_sample_id_for_result(run) for key in _get_run_output_keys(db, run): - classified = spec.classify(key, sample_id) + classified = spec.classify_output(key, sample_id) if classified: outputs[key] = classified return outputs @@ -826,7 +934,7 @@ async def list_workflow_outputs_from_s3( if not key or key in outputs: continue - classified = spec.classify(key, sample_id) + classified = spec.classify_output(key, sample_id) if classified is not None: outputs[key] = classified @@ -911,7 +1019,7 @@ async def get_result_output_downloads(db: Session, run: WorkflowRun) -> list[Res # Sort and filter outputs for key, output in sorted(outputs.items(), key=_get_output_sort_key): - if output.category == "snapshot": + if output.category in ("snapshot", "usage"): continue downloads.append( ResultDownloadItem( @@ -935,7 +1043,11 @@ async def get_all_downloads_zipped(db: Session, run: WorkflowRun) -> BytesIO: results_spec = get_output_spec(run) outputs = collect_classified_outputs(db, run, results_spec) # Exclude snapshots - downloads = [(key, output) for key, output in outputs.items() if output.category != "snapshot"] + downloads = [ + (key, output) + for key, output in outputs.items() + if output.category not in ("snapshot", "usage") + ] used_filenames: set[str] = set() zip_file = BytesIO() @@ -1028,3 +1140,14 @@ async def get_result_snapshot_downloads(db: Session, run: WorkflowRun) -> list[R ) ) return downloads + + +async def get_run_service_usage(usage_file: str) -> float | None: + content = await read_s3_file(usage_file) + csv_reader = csv.DictReader(StringIO(content)) + su_values: list[float] = [] + for row in csv_reader: + value = row.get("Service Units") + if value: + su_values.append(float(value.strip())) + return sum(su_values) diff --git a/app/services/seqera.py b/app/services/seqera.py index 53c3f54..1eec212 100644 --- a/app/services/seqera.py +++ b/app/services/seqera.py @@ -10,11 +10,14 @@ import httpx import yaml -from .seqera_client import SeqeraClient +from .seqera_client import SeqeraClient, list_workflows_raw from .seqera_errors import SeqeraAPIError, SeqeraConfigurationError logger = logging.getLogger(__name__) +# Pipeline statuses that still occupy a slot on Gadi (queued or actively running). +ACTIVE_PIPELINE_STATUSES = frozenset({"SUBMITTED", "RUNNING"}) + class WorkflowExecutorError(RuntimeError): """Raised when workflow execution through Seqera fails.""" @@ -73,6 +76,28 @@ async def post_seqera_launch( ) +async def count_active_workflows(workspace_id: str | None = None) -> int: + """Count Seqera workflow runs that are still queued or running on Gadi. + + Used to cap how many new workflows the scheduler submits at once, since each + run holds a slice of Gadi's shared PBS job-slot limit for its whole lifetime. + + Queries each active status separately via the API's `search=status:` + filter and reads the server-computed `totalSize`, rather than fetching a page + of workflows and counting client-side - `totalSize` reflects the full filtered + count regardless of the page size, so this is exact even if there are more + matching runs than fit on one page. + """ + total = 0 + for status in ACTIVE_PIPELINE_STATUSES: + data = await list_workflows_raw( + workspace_id, search_query=f"status:{status}", max_results=1 + ) + if isinstance(data, dict): + total += data.get("totalSize") or 0 + return total + + def _get_required_env(key: str) -> str: """Get required environment variable or raise error.""" value = os.getenv(key) diff --git a/app/services/seqera_client.py b/app/services/seqera_client.py index 657a4ac..9d9816b 100644 --- a/app/services/seqera_client.py +++ b/app/services/seqera_client.py @@ -85,10 +85,13 @@ def _headers(token: str) -> dict[str, str]: async def list_workflows_raw( workspace_id: str | None = None, search_query: str | None = None, + max_results: int | None = None, ) -> dict[str, Any] | list[Any]: api_url, token, params = _get_api_context(workspace_id) if search_query: params["search"] = search_query + if max_results is not None: + params["max"] = str(max_results) url = f"{api_url}/workflow" async with httpx.AsyncClient(timeout=httpx.Timeout(60)) as client: diff --git a/app/services/wisps_config.py b/app/services/wisps_config.py index e246c28..d265df5 100644 --- a/app/services/wisps_config.py +++ b/app/services/wisps_config.py @@ -5,6 +5,7 @@ from typing import Any, Literal from ..schemas.workflows import WorkflowUserDetails +from .cluster_utils import GADI_PROJECT from .workflow_config_fetcher import fetch_workflow_config WispsMode = Literal["g1-g2", "manual"] @@ -39,9 +40,7 @@ def get_wisps_config_profiles() -> list[str]: def get_wisps_config_text( config_file_path: str, *, - job_id: str, user_details: WorkflowUserDetails, - timestamp: str, ) -> str: """Read wisps config and append a process override block with runtime values. @@ -50,10 +49,6 @@ def get_wisps_config_text( """ base = fetch_workflow_config(config_file_path) - cluster_opts = ( - f"-P yz52 -v JOB_ID={job_id},USER_NAME={user_details.user_email}," - f"TIMESTAMP={timestamp},FULL_NAME={user_details.full_name}," - f"INSTITUTE={user_details.institute},IP_ADDRESS={user_details.ip_address}" - ) + cluster_opts = f"-P {GADI_PROJECT} -A {user_details.get_encoded_account_details()}" override = f'\nprocess {{\n clusterOptions = "{cluster_opts}"\n}}\n' return base + override diff --git a/app/services/wisps_executor.py b/app/services/wisps_executor.py index a0a4208..7abe379 100644 --- a/app/services/wisps_executor.py +++ b/app/services/wisps_executor.py @@ -54,7 +54,6 @@ async def prepare_wisps_workflow( raise SeqeraConfigurationError("Missing output identifier for workflow launch") out_dir = f"s3://{s3_bucket}/{output_id.strip()}" - timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") job_id = (form.runName or "").strip() if not job_id: raise SeqeraConfigurationError("Missing run name for workflow launch") @@ -72,9 +71,7 @@ async def prepare_wisps_workflow( config_text = get_wisps_config_text( config_path, - job_id=job_id, user_details=user_details, - timestamp=timestamp, ) launch_payload: dict[str, Any] = { diff --git a/app/templates/admin/system_status.html b/app/templates/admin/system_status.html index 7308a21..e790cdc 100644 --- a/app/templates/admin/system_status.html +++ b/app/templates/admin/system_status.html @@ -11,6 +11,12 @@

System Status

{% endblock %} {% block content %} +
@@ -33,7 +39,7 @@

Components

Status Latency - @@ -46,7 +52,9 @@

Components

+
+
+
+

History

+
+ + + +
+
+
+
Loading…
+
+ +
+
+ {% endblock %} diff --git a/docs/schema_diagram b/docs/schema_diagram index b554556..d078054 100644 --- a/docs/schema_diagram +++ b/docs/schema_diagram @@ -53,6 +53,9 @@ digraph { prerun_script_path TEXT () + + tool + TEXT () > URL="http://Workflow_details.html"] Workflow -> WorkflowRun [label=runs color="#1E88E5" style=dashed tooltip="Relation between Workflow and WorkflowRun"] WorkflowRun [label=< diff --git a/docs/schema_diagram.svg b/docs/schema_diagram.svg index 5f36042..8d030aa 100644 --- a/docs/schema_diagram.svg +++ b/docs/schema_diagram.svg @@ -1,62 +1,62 @@ - - - - + + + AppUser - - -AppUser + + +AppUser + + +id + + +UUID (PK) -id +auth0_user_id -UUID (PK) +TEXT (Unique) -auth0_user_id +name -TEXT (Unique) +TEXT () -name +email -TEXT () +TEXT (Unique) -email +credit -TEXT (Unique) +BIGINT () -credit +credit_updated_at -BIGINT () +DATETIME () -credit_updated_at +credit_updated_by -DATETIME () - - -credit_updated_by - - -TEXT () +TEXT () @@ -64,81 +64,81 @@ WorkflowRun - - -WorkflowRun + + +WorkflowRun + + +id + + +UUID (PK) -id +workflow_id -UUID (PK) +UUID () -workflow_id +owner_user_id -UUID () +UUID () -owner_user_id +seqera_run_id -UUID () +TEXT () -seqera_run_id +binder_name -TEXT () +TEXT () -binder_name +sample_id -TEXT () +TEXT () -sample_id +run_name -TEXT () +TEXT () -run_name +submitted_form_data -TEXT () +JSON () -submitted_form_data +work_dir -JSON () +TEXT () -work_dir +launch_ip -TEXT () +TEXT () -launch_ip +submission_timestamp -TEXT () +DATETIME () -submission_timestamp +tool -DATETIME () - - -tool - - -TEXT () +TEXT () @@ -146,71 +146,77 @@ AppUser->WorkflowRun - - + + -workflow_runs +workflow_runs WorkflowRun->AppUser - - + + -owner +owner Workflow - - -Workflow + + +Workflow + + +id + + +UUID (PK) -id +name -UUID (PK) +TEXT () -name +description -TEXT () +TEXT () -description +repo_url -TEXT () +TEXT () -repo_url +default_revision -TEXT () +TEXT () -default_revision +config_path -TEXT () +TEXT () -config_path +prerun_script_path -TEXT () +TEXT () -prerun_script_path +tool -TEXT () +TEXT () @@ -218,37 +224,37 @@ WorkflowRun->Workflow - - + + -workflow +workflow RunMetric - - -RunMetric - - -run_id - - -UUID (PK) - - -max_score - - -NUMERIC(8, 2) () - - -final_design_count - - -BIGINT () + + +RunMetric + + +run_id + + +UUID (PK) + + +max_score + + +NUMERIC(8, 2) () + + +final_design_count + + +BIGINT () @@ -256,11 +262,11 @@ WorkflowRun->RunMetric - - + + -metrics +metrics @@ -268,19 +274,19 @@ -RunInput +RunInput -run_id +run_id -UUID (PK) +UUID (PK) -s3_object_id +s3_object_id -TEXT (PK) +TEXT (PK) @@ -288,31 +294,31 @@ WorkflowRun->RunInput - + -inputs +inputs RunOutput - - -RunOutput - - -run_id - - -UUID (PK) - - -s3_object_id - - -TEXT (PK) + + +RunOutput + + +run_id + + +UUID (PK) + + +s3_object_id + + +TEXT (PK) @@ -320,41 +326,41 @@ WorkflowRun->RunOutput - - + + -outputs +outputs Workflow->WorkflowRun - - + + -runs +runs RunMetric->WorkflowRun - - + + -run +run RunInput->WorkflowRun - - + + -run +run @@ -362,31 +368,31 @@ -S3Object +S3Object -object_key +object_key -TEXT (PK) +TEXT (PK) -URI +URI -TEXT () +TEXT () -version_id +version_id -TEXT () +TEXT () -size_bytes +size_bytes -BIGINT () +BIGINT () @@ -404,17 +410,17 @@ RunOutput->WorkflowRun - - + + -run +run RunOutput->S3Object - + @@ -434,8 +440,8 @@ S3Object->RunOutput - - + + run_outputs @@ -444,75 +450,75 @@ QueuedJob - - -QueuedJob + + +QueuedJob + + +id + + +UUID (PK) -id +workflow_run_id -UUID (PK) +UUID () -workflow_run_id +workflow_id -UUID () +UUID () -workflow_id +launch_payload -UUID () +JSON () -launch_payload +status -JSON () +VARCHAR(20) () -status +attempts -VARCHAR(20) () +INTEGER () -attempts +queued_at -INTEGER () +DATETIME () -queued_at +last_attempt_at -DATETIME () +DATETIME () -last_attempt_at +next_attempt_at -DATETIME () +DATETIME () -next_attempt_at +submitted_at -DATETIME () +DATETIME () -submitted_at +error -DATETIME () - - -error - - -TEXT () +TEXT () @@ -520,59 +526,59 @@ QueuedJob->WorkflowRun - - + + -workflow_run +workflow_run QueuedJob->Workflow - - + + -workflow +workflow SystemStatusCache - - -SystemStatusCache + + +SystemStatusCache + + +key + + +TEXT (PK) -key +payload -TEXT (PK) +JSON () -payload +checked_at -JSON () +DATETIME () -checked_at +expires_at -DATETIME () +DATETIME () -expires_at +updated_at -DATETIME () - - -updated_at - - -DATETIME () +DATETIME () diff --git a/pyproject.toml b/pyproject.toml index 9ed45c7..67a813c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,6 @@ dependencies = [ "cachetools~=5.5", "python-jose[cryptography]~=3.5", "starlette-admin~=0.16", - "unidecode~=1.3", "apscheduler>=3.11.3", "loguru>=0.7.3", "typer>=0.26.8", @@ -42,6 +41,7 @@ dev = [ "pytest-cov~=4.1", "pytest-mock~=3.12", "httpx~=0.28", + "httpx2~=2.9", "coverage[toml]~=7.4", "ruff~=0.14", "black~=26.5", diff --git a/tests/conftest.py b/tests/conftest.py index f51b41d..6f99103 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -140,6 +140,7 @@ def persistent_models(test_db): datagen.WorkflowRunFactory, datagen.RunInputFactory, datagen.RunOutputFactory, + datagen.S3ObjectFactory, datagen.QueuedJobFactory, ] @@ -182,6 +183,7 @@ def app(test_engine): Workflow( id=uuid4(), name="de-novo-design", + tool="bindcraft", description="Test workflow", repo_url="https://github.com/test/repo", default_revision="dev", diff --git a/tests/datagen.py b/tests/datagen.py index 11fd9af..e9db6e4 100644 --- a/tests/datagen.py +++ b/tests/datagen.py @@ -5,7 +5,7 @@ from polyfactory.factories.dataclass_factory import DataclassFactory from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory -from app.db.models.core import AppUser, RunInput, RunOutput, Workflow, WorkflowRun +from app.db.models.core import AppUser, RunInput, RunOutput, S3Object, Workflow, WorkflowRun from app.db.models.job_queue import QueuedJob from app.services.job_utils import UserJobListRow @@ -62,6 +62,10 @@ class RunOutputFactory(SQLAlchemyFactory[RunOutput]): __set_relationships__ = False +class S3ObjectFactory(SQLAlchemyFactory[S3Object]): + __set_relationships__ = False + + class QueuedJobFactory(SQLAlchemyFactory[QueuedJob]): __set_relationships__ = False diff --git a/tests/scheduler/test_job_queue.py b/tests/scheduler/test_job_queue.py index 67ba32b..4ed0225 100644 --- a/tests/scheduler/test_job_queue.py +++ b/tests/scheduler/test_job_queue.py @@ -1,4 +1,5 @@ from datetime import UTC, datetime, timedelta +from uuid import uuid4 from apscheduler.jobstores.memory import MemoryJobStore from apscheduler.schedulers.blocking import BlockingScheduler @@ -20,11 +21,11 @@ def _get_db(): def _create_queued_job(*, status: str = "pending", next_attempt_at: datetime | None = None): user = AppUserFactory.create_sync() - workflow = WorkflowFactory.create_sync(name="de-novo-design") + workflow = WorkflowFactory.create_sync() workflow_run = WorkflowRunFactory.create_sync( owner=user, workflow=workflow, - work_dir=f"/work/{status}-{next_attempt_at.timestamp() if next_attempt_at else 'none'}", + work_dir=f"/work/{status}-{next_attempt_at.timestamp() if next_attempt_at else 'none'}-{uuid4()}", ) return QueuedJobFactory.create_sync( workflow=workflow, @@ -73,6 +74,11 @@ def test_submit_pending_jobs_schedules_only_due_pending_jobs( monkeypatch.setattr(scheduler_jobs, "get_db", _get_db_override(test_db)) monkeypatch.setattr(scheduler_jobs, "SCHEDULER", scheduler) monkeypatch.setattr(scheduler_jobs, "is_seqera_available", lambda _db_session: True) + monkeypatch.setattr( + scheduler_jobs, + "get_available_workflow_capacity", + lambda: scheduler_jobs.MAX_CONCURRENT_WORKFLOWS, + ) scheduler_jobs.submit_pending_jobs(dry_run=True) @@ -103,9 +109,71 @@ def test_submit_pending_jobs_skips_jobs_already_scheduled(test_db, persistent_mo monkeypatch.setattr(scheduler_jobs, "get_db", _get_db_override(test_db)) monkeypatch.setattr(scheduler_jobs, "SCHEDULER", scheduler) monkeypatch.setattr(scheduler_jobs, "is_seqera_available", lambda _db_session: True) + monkeypatch.setattr( + scheduler_jobs, + "get_available_workflow_capacity", + lambda: scheduler_jobs.MAX_CONCURRENT_WORKFLOWS, + ) scheduler_jobs.submit_pending_jobs() scheduled_jobs = scheduler.get_jobs(jobstore="memory") assert [job.id for job in scheduled_jobs] == [launch_id] assert scheduled_jobs[0].kwargs == {"job_id": due_job.id, "dry_run": False} + + +def test_submit_pending_jobs_skips_when_no_gadi_capacity(test_db, persistent_models, monkeypatch): + _create_queued_job(next_attempt_at=datetime.now(UTC) - timedelta(minutes=1)) + scheduler = _make_scheduler() + + monkeypatch.setattr(scheduler_jobs, "get_db", _get_db_override(test_db)) + monkeypatch.setattr(scheduler_jobs, "SCHEDULER", scheduler) + monkeypatch.setattr(scheduler_jobs, "is_seqera_available", lambda _db_session: True) + monkeypatch.setattr(scheduler_jobs, "get_available_workflow_capacity", lambda: 0) + + scheduler_jobs.submit_pending_jobs() + + assert scheduler.get_jobs(jobstore="memory") == [] + + +def test_submit_pending_jobs_caps_submissions_to_available_capacity( + test_db, persistent_models, monkeypatch +): + now = datetime.now(UTC) + due_jobs = [ + _create_queued_job(status="pending", next_attempt_at=now - timedelta(minutes=1)) + for _ in range(3) + ] + scheduler = _make_scheduler() + + monkeypatch.setattr(scheduler_jobs, "get_db", _get_db_override(test_db)) + monkeypatch.setattr(scheduler_jobs, "SCHEDULER", scheduler) + monkeypatch.setattr(scheduler_jobs, "is_seqera_available", lambda _db_session: True) + monkeypatch.setattr(scheduler_jobs, "get_available_workflow_capacity", lambda: 2) + + scheduler_jobs.submit_pending_jobs() + + scheduled_jobs = scheduler.get_jobs(jobstore="memory") + assert len(scheduled_jobs) == 2 + scheduled_job_ids = {job.kwargs["job_id"] for job in scheduled_jobs} + assert scheduled_job_ids.issubset({job.id for job in due_jobs}) + + +def test_get_available_workflow_capacity_uses_seqera_active_count(monkeypatch): + async def _count_active_workflows(): + return 20 + + monkeypatch.setattr(scheduler_jobs.seqera, "count_active_workflows", _count_active_workflows) + monkeypatch.setattr(scheduler_jobs, "MAX_CONCURRENT_WORKFLOWS", 25) + + assert scheduler_jobs.get_available_workflow_capacity() == 5 + + +def test_get_available_workflow_capacity_floors_at_zero(monkeypatch): + async def _count_active_workflows(): + return 30 + + monkeypatch.setattr(scheduler_jobs.seqera, "count_active_workflows", _count_active_workflows) + monkeypatch.setattr(scheduler_jobs, "MAX_CONCURRENT_WORKFLOWS", 25) + + assert scheduler_jobs.get_available_workflow_capacity() == 0 diff --git a/tests/scheduler/test_scheduler_jobs.py b/tests/scheduler/test_scheduler_jobs.py index c6747b0..ed800e3 100644 --- a/tests/scheduler/test_scheduler_jobs.py +++ b/tests/scheduler/test_scheduler_jobs.py @@ -20,7 +20,9 @@ async def _failing_launch(**_kwargs): raise RuntimeError("Seqera launch failed") -def _create_queued_job(*, attempts: int = 0, workflow_name: str = "de-novo-design"): +def _create_queued_job( + *, attempts: int = 0, workflow_name: str = "de-novo-design", tool: str | None = None +): user = AppUserFactory.create_sync() workflow = WorkflowFactory.create_sync(name=workflow_name) workflow_run = WorkflowRunFactory.create_sync( @@ -28,6 +30,7 @@ def _create_queued_job(*, attempts: int = 0, workflow_name: str = "de-novo-desig workflow=workflow, work_dir=f"/work/{workflow_name}-{attempts}-{uuid4()}", seqera_run_id=None, + tool=tool, ) return QueuedJobFactory.create_sync( workflow_run=workflow_run, @@ -125,6 +128,56 @@ async def _successful_launch(**kwargs): assert queued_job.next_attempt_at is not None +def test_launch_job_dispatches_proteindj_for_rfdiffusion_tool( + test_db, persistent_models, monkeypatch +): + queued_job = _create_queued_job(tool="rfdiffusion") + calls = [] + + async def _successful_launch(**kwargs): + calls.append(kwargs) + return WorkflowLaunchResult(workflow_id="seqera-run-rfd", status="submitted") + + async def _unexpected_bindflow_launch(**_kwargs): + raise AssertionError("launch_bindflow_workflow should not be called for rfdiffusion") + + monkeypatch.setattr(scheduler_jobs, "get_db", _get_db_override(test_db)) + monkeypatch.setattr(scheduler_jobs, "is_seqera_available", lambda _db: True) + monkeypatch.setattr(scheduler_jobs, "launch_proteindj_workflow", _successful_launch) + monkeypatch.setattr(scheduler_jobs, "launch_bindflow_workflow", _unexpected_bindflow_launch) + + scheduler_jobs.launch_job(queued_job.id) + + test_db.refresh(queued_job) + assert calls == [{"queued_job": queued_job, "dry_run": False}] + assert queued_job.workflow_run.seqera_run_id == "seqera-run-rfd" + assert queued_job.status == "submitted" + + +def test_launch_job_dispatches_bindflow_for_bindcraft_tool(test_db, persistent_models, monkeypatch): + queued_job = _create_queued_job(tool="bindcraft") + calls = [] + + async def _successful_launch(**kwargs): + calls.append(kwargs) + return WorkflowLaunchResult(workflow_id="seqera-run-bc", status="submitted") + + async def _unexpected_proteindj_launch(**_kwargs): + raise AssertionError("launch_proteindj_workflow should not be called for bindcraft") + + monkeypatch.setattr(scheduler_jobs, "get_db", _get_db_override(test_db)) + monkeypatch.setattr(scheduler_jobs, "is_seqera_available", lambda _db: True) + monkeypatch.setattr(scheduler_jobs, "launch_bindflow_workflow", _successful_launch) + monkeypatch.setattr(scheduler_jobs, "launch_proteindj_workflow", _unexpected_proteindj_launch) + + scheduler_jobs.launch_job(queued_job.id) + + test_db.refresh(queued_job) + assert calls == [{"queued_job": queued_job, "dry_run": False}] + assert queued_job.workflow_run.seqera_run_id == "seqera-run-bc" + assert queued_job.status == "submitted" + + def test_launch_job_dispatches_wisps_workflows(test_db, persistent_models, monkeypatch): queued_job = _create_queued_job(workflow_name="interaction-screening") calls = [] diff --git a/tests/test_proteindj_coverage.py b/tests/test_proteindj_coverage.py new file mode 100644 index 0000000..7afd0b9 --- /dev/null +++ b/tests/test_proteindj_coverage.py @@ -0,0 +1,603 @@ +"""Tests to boost coverage for proteindj executor and config modules.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import AsyncMock, Mock, mock_open, patch + +import pytest +from sqlalchemy import select + +from app.db.models import QueuedJob +from app.schemas.workflows import ( + ProteinDjFormData, + WorkflowFormData, + WorkflowLaunchForm, + WorkflowUserDetails, +) +from app.services.proteindj_config import ( + get_proteindj_config_profiles, + get_proteindj_config_text, + get_proteindj_default_params, +) +from app.services.proteindj_executor import ( + _design_length, + launch_proteindj_workflow, + prepare_proteindj_workflow, +) +from app.services.seqera import WorkflowLaunchResult +from app.services.seqera_errors import SeqeraConfigurationError +from tests.datagen import AppUserFactory, QueuedJobFactory, WorkflowFactory, WorkflowRunFactory + +_USER_DETAILS = WorkflowUserDetails( + user_email="user@ex.com", + ip_address="1.2.3.4", +) + + +def _form_data(**extra) -> WorkflowFormData: + return WorkflowFormData(workflow="de-novo-design", tool="rfdiffusion", **extra) + + +def _make_launch_form(**kwargs) -> WorkflowLaunchForm: + defaults = { + "workflow": "de-novo-design", + "tool": "rfdiffusion", + "runName": "test-run", + "paramsText": None, + } + defaults.update(kwargs) + return WorkflowLaunchForm(**defaults) + + +def _queued_proteindj_job( + *, + params_text: str | None = None, + prerun_script_path: str | None = None, +) -> QueuedJob: + user = AppUserFactory.create_sync() + workflow = WorkflowFactory.create_sync( + name="de-novo-design", + prerun_script_path=prerun_script_path, + ) + workflow_run = WorkflowRunFactory.create_sync(workflow=workflow, owner=user) + launch_payload = { + "computeEnvId": "ce_456", + "runName": "test-run", + "pipeline": "https://github.com/org/proteindj", + "workDir": "/work/dir", + "workspaceId": "ws_123", + "revision": "dev", + "paramsText": params_text + or ("out_dir: s3://my-bucket/run-output-id\ninput_pdb: s3://my-bucket/inputs/test.pdb"), + "configProfiles": ["singularity"], + "configText": "config_text", + "resume": False, + } + return QueuedJobFactory.create_sync( + workflow=workflow, + workflow_run=workflow_run, + launch_payload=launch_payload, + status="pending", + ) + + +@contextmanager +def _mock_proteindj_db_context(): + workflow = Mock(name="workflow") + workflow_run = Mock(name="workflow_run") + workflow_run.workflow = workflow + db_session = Mock(name="db_session") + queued_job = Mock(name="queued_job") + with patch( + "app.services.proteindj_executor.QueuedJob", return_value=queued_job + ) as queued_job_cls: + yield db_session, workflow_run, workflow, queued_job_cls, queued_job + + +@pytest.fixture +def seqera_env(monkeypatch): + """Set required Seqera environment variables for launch tests.""" + monkeypatch.setenv("SEQERA_API_URL", "https://api.seqera.test") + monkeypatch.setenv("SEQERA_ACCESS_TOKEN", "test_token") + monkeypatch.setenv("WORK_SPACE", "ws_123") + monkeypatch.setenv("COMPUTE_ID", "ce_456") + monkeypatch.setenv("WORK_DIR", "/work/dir") + monkeypatch.setenv("AWS_S3_BUCKET", "my-bucket") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test_key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test_secret") + + +def _proteindj_form_kwargs(**overrides) -> dict: + defaults = { + "workflow": "de-novo-design", + "tool": "rfdiffusion", + "starting_pdb": "s3://bucket/in.pdb", + "target_hotspot_residues": "A20,A21", + "number_of_final_designs": 5, + "min_length": 100, + "max_length": 150, + } + defaults.update(overrides) + return defaults + + +# ============================================================================= +# Tests for ProteinDjFormData validation +# ============================================================================= + + +def test_proteindj_form_data_accepts_valid_fields(): + fields = ProteinDjFormData(**_proteindj_form_kwargs()) + assert fields.min_length == 100 + assert fields.max_length == 150 + + +def test_proteindj_form_data_rejects_too_many_hotspot_residues(): + with pytest.raises(ValueError, match="Too many hotspot residues"): + ProteinDjFormData( + **_proteindj_form_kwargs(target_hotspot_residues="A1,A2,A3,A4,A5,A6,A7,A8,A9") + ) + + +def test_proteindj_form_data_allows_exactly_max_hotspot_residues(): + fields = ProteinDjFormData( + **_proteindj_form_kwargs(target_hotspot_residues="A1,A2,A3,A4,A5,A6,A7,A8") + ) + assert fields.target_hotspot_residues == "A1,A2,A3,A4,A5,A6,A7,A8" + + +def test_proteindj_form_data_rejects_min_length_below_floor(): + with pytest.raises(ValueError, match="min_length"): + ProteinDjFormData(**_proteindj_form_kwargs(min_length=64, max_length=150)) + + +def test_proteindj_form_data_rejects_max_length_above_ceiling(): + with pytest.raises(ValueError, match="max_length"): + ProteinDjFormData(**_proteindj_form_kwargs(min_length=65, max_length=151)) + + +def test_proteindj_form_data_rejects_min_length_above_max_length(): + with pytest.raises(ValueError, match="min_length must not exceed max_length"): + ProteinDjFormData(**_proteindj_form_kwargs(min_length=150, max_length=100)) + + +def test_proteindj_form_data_allows_boundary_lengths(): + fields = ProteinDjFormData(**_proteindj_form_kwargs(min_length=65, max_length=150)) + assert fields.min_length == 65 + assert fields.max_length == 150 + + +# ============================================================================= +# Tests for get_proteindj_default_params() +# ============================================================================= + + +def test_get_proteindj_default_params_all_fields(): + params = get_proteindj_default_params( + "s3://bucket/out", + input_pdb="s3://bucket/in.pdb", + hotspot_residues="A20,A21", + num_designs=5, + design_length="100-150", + ) + assert params == { + "out_dir": "s3://bucket/out", + "input_pdb": "s3://bucket/in.pdb", + "hotspot_residues": "A20,A21", + "num_designs": 5, + "design_length": "100-150", + } + + +def test_get_proteindj_default_params_missing_required_field_raises(): + with pytest.raises(TypeError): + get_proteindj_default_params("s3://bucket/out", num_designs=3) # type: ignore[call-arg] + + +# ============================================================================= +# Tests for get_proteindj_config_profiles() +# ============================================================================= + + +def test_get_proteindj_config_profiles_returns_list(): + profiles = get_proteindj_config_profiles() + assert isinstance(profiles, list) + + +def test_get_proteindj_config_profiles_contains_singularity(): + assert "singularity" in get_proteindj_config_profiles() + + +# ============================================================================= +# Tests for get_proteindj_config_text() +# ============================================================================= + + +def test_get_proteindj_config_text_appends_process_block(): + with patch("builtins.open", mock_open(read_data="base_config")): + result = get_proteindj_config_text("/fake/proteindj.config", user_details=_USER_DETAILS) + assert "process {" in result + assert "clusterOptions" in result + + +def test_get_proteindj_config_text_contains_encoded_email_and_encoded_ip(): + with patch("builtins.open", mock_open(read_data="base_config")): + result = get_proteindj_config_text("/fake/proteindj.config", user_details=_USER_DETAILS) + assert "dXNlckBleC5jb20=" in result + assert "MS4yLjMuNA==" in result + + +def test_get_proteindj_config_text_without_ip_uses_encoded_email_only(): + with patch("builtins.open", mock_open(read_data="base_config")): + result = get_proteindj_config_text( + "/fake/proteindj.config", + user_details=_USER_DETAILS.model_copy(update={"ip_address": ""}), + ) + assert "-A dXNlckBleC5jb20=" in result + assert ":" not in result.split("clusterOptions = ")[1] + + +def test_get_proteindj_config_text_contains_base_config(): + with patch("builtins.open", mock_open(read_data="base_config")): + result = get_proteindj_config_text("/fake/proteindj.config", user_details=_USER_DETAILS) + assert "base_config" in result + + +# ============================================================================= +# Tests for _design_length() +# ============================================================================= + + +def test_design_length_formats_min_max_range(): + fields = ProteinDjFormData( + workflow="de-novo-design", + tool="rfdiffusion", + starting_pdb="s3://bucket/in.pdb", + target_hotspot_residues="A20", + number_of_final_designs=5, + min_length=100, + max_length=150, + ) + assert _design_length(fields) == "100-150" + + +# ============================================================================= +# Tests for prepare_proteindj_workflow() +# ============================================================================= + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_writes_expected_queued_job( + test_db, persistent_models, seqera_env +): + user = AppUserFactory.create_sync() + workflow = WorkflowFactory.create_sync(name="de-novo-design") + workflow_run = WorkflowRunFactory.create_sync(workflow=workflow, owner=user) + + form = _make_launch_form(runName="queued-proteindj-run") + form_data = _form_data( + starting_pdb="s3://my-bucket/inputs/test.pdb", + target_hotspot_residues="A20,A21", + number_of_final_designs=5, + min_length=100, + max_length=150, + ) + + with ( + patch( + "app.services.proteindj_executor.get_proteindj_config_text", + return_value="config_text", + ), + patch( + "app.services.proteindj_executor.get_proteindj_config_profiles", + return_value=["singularity"], + ), + ): + prepared_job = await prepare_proteindj_workflow( + form=form, + db_session=test_db, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + revision="main", + output_id="run-output-id", + form_data=form_data, + user_details=_USER_DETAILS.model_copy( + update={"user_email": "test@example.com", "ip_address": "127.0.0.1"} + ), + ) + + queued_job = test_db.scalar( + select(QueuedJob).where(QueuedJob.workflow_run_id == workflow_run.id) + ) + assert queued_job is not None + assert queued_job.workflow_id == workflow.id + assert queued_job.workflow_run_id == workflow_run.id + assert queued_job.status == "pending" + assert queued_job.next_attempt_at is not None + assert queued_job.id == prepared_job.id + assert queued_job.launch_payload["computeEnvId"] == "ce_456" + assert queued_job.launch_payload["runName"] == "queued-proteindj-run" + assert queued_job.launch_payload["pipeline"] == "https://github.com/org/proteindj" + assert queued_job.launch_payload["workDir"] == "/work/dir" + assert queued_job.launch_payload["workspaceId"] == "ws_123" + assert queued_job.launch_payload["revision"] == "main" + assert queued_job.launch_payload["configProfiles"] == ["singularity"] + assert queued_job.launch_payload["configText"] == "config_text" + assert "preRunScript" not in queued_job.launch_payload + assert queued_job.launch_payload["resume"] is False + params_text = queued_job.launch_payload["paramsText"] + assert "out_dir: s3://my-bucket/run-output-id" in params_text + assert "input_pdb: s3://my-bucket/inputs/test.pdb" in params_text + assert "hotspot_residues: A20,A21" in params_text + assert "num_designs: 5" in params_text + assert "design_length: 100-150" in params_text + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_appends_custom_params_text( + test_db, persistent_models, seqera_env +): + user = AppUserFactory.create_sync() + workflow = WorkflowFactory.create_sync(name="de-novo-design") + workflow_run = WorkflowRunFactory.create_sync(workflow=workflow, owner=user) + + form = _make_launch_form(paramsText="extra_param: value") + + with ( + patch("app.services.proteindj_executor.get_proteindj_config_text", return_value=""), + patch( + "app.services.proteindj_executor.get_proteindj_config_profiles", + return_value=["singularity"], + ), + ): + prepared_job = await prepare_proteindj_workflow( + form=form, + db_session=test_db, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id="run-output-id", + form_data=_form_data( + starting_pdb="s3://my-bucket/inputs/test.pdb", + target_hotspot_residues="A20,A21", + number_of_final_designs=5, + min_length=100, + max_length=150, + ), + user_details=_USER_DETAILS, + ) + + queued_job = test_db.scalar( + select(QueuedJob).where(QueuedJob.workflow_run_id == workflow_run.id) + ) + assert "extra_param: value" in queued_job.launch_payload["paramsText"] + assert prepared_job.id == queued_job.id + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_missing_run_name(seqera_env): + form = _make_launch_form(runName=" ") + with ( + _mock_proteindj_db_context() as (db_session, workflow_run, *_), + pytest.raises(SeqeraConfigurationError, match="run name"), + ): + await prepare_proteindj_workflow( + form=form, + db_session=db_session, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id="run-output-id", + form_data=_form_data(), + user_details=_USER_DETAILS, + ) + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_missing_output_id(seqera_env): + form = _make_launch_form() + with ( + _mock_proteindj_db_context() as (db_session, workflow_run, *_), + pytest.raises(SeqeraConfigurationError, match="output identifier"), + ): + await prepare_proteindj_workflow( + form=form, + db_session=db_session, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id=None, + form_data=_form_data(), + user_details=_USER_DETAILS, + ) + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_empty_output_id(seqera_env): + form = _make_launch_form() + with ( + _mock_proteindj_db_context() as (db_session, workflow_run, *_), + pytest.raises(SeqeraConfigurationError, match="output identifier"), + ): + await prepare_proteindj_workflow( + form=form, + db_session=db_session, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id=" ", + form_data=_form_data(), + user_details=_USER_DETAILS, + ) + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_missing_starting_pdb(seqera_env): + form = _make_launch_form() + with ( + _mock_proteindj_db_context() as (db_session, workflow_run, *_), + pytest.raises(SeqeraConfigurationError, match="starting_pdb"), + ): + await prepare_proteindj_workflow( + form=form, + db_session=db_session, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id="run-output-id", + form_data=_form_data( + target_hotspot_residues="A20,A21", + number_of_final_designs=5, + min_length=100, + max_length=150, + ), + user_details=_USER_DETAILS, + ) + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_missing_hotspot_residues(seqera_env): + form = _make_launch_form() + with ( + _mock_proteindj_db_context() as (db_session, workflow_run, *_), + pytest.raises(SeqeraConfigurationError, match="target_hotspot_residues"), + ): + await prepare_proteindj_workflow( + form=form, + db_session=db_session, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id="run-output-id", + form_data=_form_data( + starting_pdb="s3://my-bucket/inputs/test.pdb", + number_of_final_designs=5, + min_length=100, + max_length=150, + ), + user_details=_USER_DETAILS, + ) + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_missing_num_designs(seqera_env): + form = _make_launch_form() + with ( + _mock_proteindj_db_context() as (db_session, workflow_run, *_), + pytest.raises(SeqeraConfigurationError, match="number_of_final_designs"), + ): + await prepare_proteindj_workflow( + form=form, + db_session=db_session, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id="run-output-id", + form_data=_form_data( + starting_pdb="s3://my-bucket/inputs/test.pdb", + target_hotspot_residues="A20,A21", + min_length=100, + max_length=150, + ), + user_details=_USER_DETAILS, + ) + + +@pytest.mark.anyio +async def test_prepare_proteindj_workflow_missing_design_length(seqera_env): + form = _make_launch_form() + with ( + _mock_proteindj_db_context() as (db_session, workflow_run, *_), + pytest.raises(SeqeraConfigurationError, match="min_length"), + ): + await prepare_proteindj_workflow( + form=form, + db_session=db_session, + workflow_run=workflow_run, + pipeline="https://github.com/org/proteindj", + config_path="/fake/proteindj.config", + output_id="run-output-id", + form_data=_form_data( + starting_pdb="s3://my-bucket/inputs/test.pdb", + target_hotspot_residues="A20,A21", + number_of_final_designs=5, + ), + user_details=_USER_DETAILS, + ) + + +# ============================================================================= +# Tests for launch_proteindj_workflow() +# ============================================================================= + + +@pytest.mark.anyio +async def test_launch_proteindj_workflow_success(seqera_env, persistent_models): + expected_result = WorkflowLaunchResult( + workflow_id="wf_success", status="submitted", message=None + ) + + with ( + patch( + "app.services.proteindj_executor.post_seqera_launch", + new_callable=AsyncMock, + return_value=expected_result, + ) as mock_post, + ): + result = await launch_proteindj_workflow(queued_job=_queued_proteindj_job()) + + assert result.workflow_id == "wf_success" + assert result.status == "submitted" + mock_post.assert_called_once() + posted_payload = mock_post.call_args.args[0]["launch"] + assert "module load singularity" in posted_payload["preRunScript"] + assert "module load nextflow" in posted_payload["preRunScript"] + assert "export AWS_ACCESS_KEY_ID" in posted_payload["preRunScript"] + + +@pytest.mark.anyio +async def test_launch_proteindj_workflow_with_prerun_script_path(seqera_env, persistent_models): + expected_result = WorkflowLaunchResult(workflow_id="wf_prerun", status="submitted") + + with ( + patch( + "app.services.proteindj_executor.post_seqera_launch", + new_callable=AsyncMock, + return_value=expected_result, + ) as mock_post, + patch( + "app.services.proteindj_executor.get_executor_script", + return_value="prerun_body", + ) as mock_script, + ): + result = await launch_proteindj_workflow( + queued_job=_queued_proteindj_job(prerun_script_path="/some/prerun.sh") + ) + + assert result.workflow_id == "wf_prerun" + posted_payload = mock_post.call_args.args[0]["launch"] + assert posted_payload["preRunScript"] == "prerun_body" + assert mock_script.call_args.kwargs["prerun_script_path"] == "/some/prerun.sh" + + +@pytest.mark.anyio +async def test_launch_proteindj_workflow_dry_run(seqera_env, persistent_models): + with patch( + "app.services.proteindj_executor.post_seqera_launch", + new_callable=AsyncMock, + ) as mock_post: + result = await launch_proteindj_workflow(queued_job=_queued_proteindj_job(), dry_run=True) + + assert result is None + mock_post.assert_not_called() + + +@pytest.mark.anyio +async def test_launch_proteindj_workflow_missing_env_var(monkeypatch, persistent_models): + monkeypatch.delenv("SEQERA_API_URL", raising=False) + monkeypatch.delenv("SEQERA_ACCESS_TOKEN", raising=False) + + with pytest.raises(SeqeraConfigurationError, match="SEQERA_API_URL"): + await launch_proteindj_workflow(queued_job=_queued_proteindj_job()) diff --git a/tests/test_proteinfold_coverage.py b/tests/test_proteinfold_coverage.py index ecd9d80..3cb71a8 100644 --- a/tests/test_proteinfold_coverage.py +++ b/tests/test_proteinfold_coverage.py @@ -35,8 +35,6 @@ _USER_DETAILS = WorkflowUserDetails( user_email="user@ex.com", - full_name="Test_User", - institute="USYD", ip_address="1.2.3.4", ) @@ -143,8 +141,8 @@ def test_tool_params_with_int(): def test_tool_params_with_str(): - result = _tool_params(_form_data(alphafold2_random_seed="42")) - assert result == {"alphafold2_random_seed": "42"} + result = _tool_params(_form_data(random_seed="42")) + assert result == {"random_seed": "42"} def test_tool_params_none_value_excluded(): @@ -369,7 +367,6 @@ async def test_prepare_proteinfold_workflow_writes_expected_queued_job( user_details=_USER_DETAILS.model_copy( update={ "user_email": "test@example.com", - "institute": "example.com", "ip_address": "127.0.0.1", } ), @@ -544,33 +541,36 @@ def test_get_proteinfold_config_text_appends_process_block(): with patch("builtins.open", mock_open(read_data="base_config")): result = get_proteinfold_config_text( "/fake/proteinfold.config", - job_id="my-job", user_details=_USER_DETAILS, - timestamp="20240101_120000", ) assert "process {" in result assert "clusterOptions" in result -def test_get_proteinfold_config_text_contains_job_fields(): +def test_get_proteinfold_config_text_contains_encoded_email_and_encoded_ip(): with patch("builtins.open", mock_open(read_data="base_config")): result = get_proteinfold_config_text( "/fake/proteinfold.config", - job_id="my-job", user_details=_USER_DETAILS, - timestamp="20240101_120000", ) - assert "my-job" in result - assert "user@ex.com" in result - assert "20240101_120000" in result + assert "dXNlckBleC5jb20=" in result + assert "MS4yLjMuNA==" in result + + +def test_get_proteinfold_config_text_without_ip_uses_encoded_email_only(): + with patch("builtins.open", mock_open(read_data="base_config")): + result = get_proteinfold_config_text( + "/fake/proteinfold.config", + user_details=_USER_DETAILS.model_copy(update={"ip_address": ""}), + ) + assert "-A dXNlckBleC5jb20=" in result + assert ":" not in result.split("clusterOptions = ")[1] def test_get_proteinfold_config_text_contains_base_config(): with patch("builtins.open", mock_open(read_data="base_config")): result = get_proteinfold_config_text( "/fake/proteinfold.config", - job_id="my-job", user_details=_USER_DETAILS, - timestamp="20240101_120000", ) assert "base_config" in result diff --git a/tests/test_routes_admin_system_status.py b/tests/test_routes_admin_system_status.py index 1834af6..576310f 100644 --- a/tests/test_routes_admin_system_status.py +++ b/tests/test_routes_admin_system_status.py @@ -3,13 +3,15 @@ from __future__ import annotations import os -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from pytest import MonkeyPatch from app.db.admin import require_admin_access +from app.db.models.system_status import SystemStatusIncident from app.routes.system_status import router as system_status_router from app.schemas.health import ProbeResult, SystemStatus from app.services import health @@ -136,3 +138,97 @@ def test_system_status_not_shadowed_by_admin_mount(mocker): assert resp.status_code == 401 # our endpoint's auth gate, not the admin mount assert "text/html" not in resp.headers.get("content-type", "") + + +# --------------------------------------------------------------------------- +# GET /admin/api/system-status/history (downtime incidents) +# --------------------------------------------------------------------------- + + +def _seed_incident(test_db, *, component, status, hours_ago, ended_hours_ago=None, message=None): + test_db.add( + SystemStatusIncident( + component=component, + status=status, + started_at=datetime.now(UTC) - timedelta(hours=hours_ago), + ended_at=( + datetime.now(UTC) - timedelta(hours=ended_hours_ago) + if ended_hours_ago is not None + else None + ), + message=message, + ) + ) + test_db.commit() + + +def test_history_requires_admin(client): + resp = client.get("/admin/api/system-status/history") + assert resp.status_code == 401 + + +def test_history_returns_per_component_summary_and_incidents(client, test_db): + client.app.dependency_overrides[require_admin_access] = lambda: {"sub": "auth0|admin"} + _seed_incident( + test_db, + component="seqera_compute_env", + status="unhealthy", + hours_ago=2, + ended_hours_ago=1, + message="ERRORED", + ) + + resp = client.get("/admin/api/system-status/history?hours=24") + assert resp.status_code == 200 + body = resp.json() + assert body["windowHours"] == 24 + assert "since" in body and "until" in body + + by_name = {c["name"]: c for c in body["components"]} + ce = by_name["seqera_compute_env"] + assert ce["summary"]["incidentCount"] == 1 + assert ce["summary"]["downtimeSeconds"] == pytest.approx(3600, rel=0.01) + assert ce["incidents"][0]["message"] == "ERRORED" + assert ce["incidents"][0]["endedAt"] is not None + + +def test_history_includes_ongoing_incident(client, test_db): + client.app.dependency_overrides[require_admin_access] = lambda: {"sub": "auth0|admin"} + _seed_incident(test_db, component="seqera_api", status="unhealthy", hours_ago=1) + + resp = client.get("/admin/api/system-status/history?hours=24") + assert resp.status_code == 200 + body = resp.json() + api = next(c for c in body["components"] if c["name"] == "seqera_api") + assert api["incidents"][0]["endedAt"] is None + assert api["summary"]["incidentCount"] == 1 + + +def test_history_excludes_incidents_outside_window(client, test_db): + client.app.dependency_overrides[require_admin_access] = lambda: {"sub": "auth0|admin"} + _seed_incident( + test_db, component="seqera_api", status="unhealthy", hours_ago=48, ended_hours_ago=47 + ) + + resp = client.get("/admin/api/system-status/history?hours=24") + assert resp.status_code == 200 + assert resp.json()["components"] == [] + + +def test_history_filters_by_component_query_param(client, test_db): + client.app.dependency_overrides[require_admin_access] = lambda: {"sub": "auth0|admin"} + _seed_incident(test_db, component="seqera_api", status="unhealthy", hours_ago=1) + _seed_incident(test_db, component="seqera_compute_env", status="unhealthy", hours_ago=1) + + resp = client.get("/admin/api/system-status/history?hours=24&component=seqera_api") + assert resp.status_code == 200 + body = resp.json() + assert len(body["components"]) == 1 + assert body["components"][0]["name"] == "seqera_api" + + +def test_history_empty_when_no_incidents_recorded(client): + client.app.dependency_overrides[require_admin_access] = lambda: {"sub": "auth0|admin"} + resp = client.get("/admin/api/system-status/history") + assert resp.status_code == 200 + assert resp.json()["components"] == [] diff --git a/tests/test_routes_results.py b/tests/test_routes_results.py index cd63741..18fbf8d 100644 --- a/tests/test_routes_results.py +++ b/tests/test_routes_results.py @@ -24,7 +24,12 @@ def _configure_bindcraft_run(run: WorkflowRun) -> None: - run.workflow = Workflow(name="de-novo-design") + run.workflow = Workflow( + name="de-novo-design", + repo_url="https://github.com/test/de-novo-design", + default_revision="main", + config_path="/config/de-novo-design.config", + ) run.submitted_form_data = {"mode": "bindcraft"} @@ -332,7 +337,12 @@ async def test_get_result_downloads_returns_presigned_links_for_tracked_outputs( name="Results User 5", email="results5@example.com", ) - workflow = Workflow(name="de-novo-design") + workflow = Workflow( + name="de-novo-design", + repo_url="https://github.com/test/de-novo-design", + default_revision="main", + config_path="/config/de-novo-design.config", + ) run = WorkflowRun( owner=user, workflow=workflow, @@ -492,7 +502,12 @@ async def test_get_result_downloads_returns_presigned_links_for_proteinfold_outp name=f"Proteinfold Downloads {tool}", email=f"results-proteinfold-downloads-{tool}@example.com", ) - workflow = Workflow(name="single-prediction") + workflow = Workflow( + name="single-prediction", + repo_url="https://github.com/test/single-prediction", + default_revision="main", + config_path="/config/single-prediction.config", + ) run = WorkflowRun( owner=user, workflow=workflow, @@ -622,7 +637,12 @@ async def test_get_result_snapshots_returns_presigned_links_for_tracked_outputs( name="Results User Snapshots 1", email="results-snapshots1@example.com", ) - workflow = Workflow(name="de-novo-design") + workflow = Workflow( + name="de-novo-design", + repo_url="https://github.com/test/de-novo-design", + default_revision="main", + config_path="/config/de-novo-design.config", + ) run = WorkflowRun( workflow=workflow, owner=user, @@ -973,7 +993,12 @@ async def test_get_result_setting_params_overlays_queued_job_payload(test_db): name="Queued Job User", email="queued-job@example.com", ) - workflow = Workflow(name="de-novo-design-queued") + workflow = Workflow( + name="de-novo-design-queued", + repo_url="https://github.com/test/de-novo-design-queued", + default_revision="main", + config_path="/config/de-novo-design-queued.config", + ) run = WorkflowRun( owner=user, workflow=workflow, @@ -1011,7 +1036,12 @@ async def test_get_result_setting_params_queued_job_invalid_yaml_kept_as_string( name="Queued Job User 2", email="queued-job2@example.com", ) - workflow = Workflow(name="de-novo-design-queued-2") + workflow = Workflow( + name="de-novo-design-queued-2", + repo_url="https://github.com/test/de-novo-design-queued-2", + default_revision="main", + config_path="/config/de-novo-design-queued-2.config", + ) run = WorkflowRun( owner=user, workflow=workflow, diff --git a/tests/test_routes_workflow_jobs_extra.py b/tests/test_routes_workflow_jobs_extra.py index 87d3633..4ed1537 100644 --- a/tests/test_routes_workflow_jobs_extra.py +++ b/tests/test_routes_workflow_jobs_extra.py @@ -94,7 +94,14 @@ async def test_get_job_details_success(test_db): name="User", email="user@example.com", ) - workflow = Workflow(id=uuid4(), name="BindCraft", description="Binding workflow") + workflow = Workflow( + id=uuid4(), + name="de-novo-design", + description="Binding workflow", + repo_url="https://github.com/test/bindcraft", + default_revision="main", + config_path="/config/bindcraft.config", + ) run = WorkflowRun( id=uuid4(), owner_user_id=user.id, @@ -129,7 +136,7 @@ async def test_get_job_details_success(test_db): assert result.id == str(run.id) assert result.jobName == "PDL1" - assert result.workflow == "Bindcraft" + assert result.workflow == "De Novo Design" assert result.score == 0.912 diff --git a/tests/test_routes_workflows.py b/tests/test_routes_workflows.py index 187bef1..935e231 100644 --- a/tests/test_routes_workflows.py +++ b/tests/test_routes_workflows.py @@ -33,6 +33,22 @@ async def _queue_job_for_route_prepare(form, _s3_input_key, **kwargs): return queued_job +async def _queue_job_for_proteindj_route_prepare(form, **kwargs): + # proteindj has no samplesheet, so prepare_proteindj_workflow takes no + # s3_input_key positional arg (unlike prepare_bindflow_workflow above). + db_session = kwargs["db_session"] + workflow_run = kwargs["workflow_run"] + queued_job = QueuedJob( + workflow=workflow_run.workflow, + workflow_run=workflow_run, + launch_payload={"runName": form.runName}, + status="pending", + ) + db_session.add(queued_job) + db_session.flush() + return queued_job + + @pytest.fixture def role_check_client(test_engine): """Test client with auth bypassed but require_workflow_execution_role active.""" @@ -60,6 +76,7 @@ def role_check_client(test_engine): Workflow( id=uuid4(), name="de-novo-design", + tool="bindcraft", description="Test workflow", repo_url="https://github.com/test/repo", default_revision="dev", @@ -202,6 +219,81 @@ def test_launch_queue_preparation_error(mock_prepare, client: TestClient, test_e assert count == 0 +def test_launch_de_novo_design_tool_mismatch_returns_500(client: TestClient): + """de-novo-design is matched on tool; an unconfigured tool for it is a 500.""" + payload = { + "launch": { + "workflow": "de-novo-design", + "tool": "rfdiffusion", + "runName": "test-run", + }, + "s3InputKey": "inputs/samplesheets/test.csv", + "formData": {"workflow": "de-novo-design", "tool": "rfdiffusion"}, + } + + response = client.post("/api/workflows/launch", json=payload) + + assert response.status_code == 500 + assert "rfdiffusion" in response.json()["detail"] + + +def _add_rfdiffusion_workflow(test_engine): + """Helper to add a de-novo-design/rfdiffusion workflow row to the test DB.""" + with Session(test_engine) as db: + existing = db.scalar( + select(Workflow).where( + Workflow.name == "de-novo-design", Workflow.tool == "rfdiffusion" + ) + ) + if not existing: + db.add( + Workflow( + id=uuid4(), + name="de-novo-design", + tool="rfdiffusion", + description="ProteinDJ workflow", + repo_url="https://github.com/test/proteindj", + default_revision="dev", + config_path="/some/proteindj.config", + ) + ) + db.commit() + + +@patch("app.routes.workflows.prepare_bindflow_workflow") +@patch( + "app.routes.workflows.prepare_proteindj_workflow", + side_effect=_queue_job_for_proteindj_route_prepare, +) +def test_launch_de_novo_design_rfdiffusion_routes_to_proteindj( + mock_prepare_proteindj, mock_prepare_bindflow, client: TestClient, test_engine +): + """tool='rfdiffusion' on de-novo-design must dispatch to the proteindj executor.""" + _add_rfdiffusion_workflow(test_engine) + + payload = { + "launch": { + "workflow": "de-novo-design", + "tool": "rfdiffusion", + "runName": "rfd-run-1", + }, + "s3InputKey": "inputs/pdb/target.pdb", + "formData": {"workflow": "de-novo-design", "tool": "rfdiffusion"}, + } + + response = client.post("/api/workflows/launch", json=payload) + + assert response.status_code == 201 + data = response.json() + assert data["status"] == "pending" + mock_prepare_proteindj.assert_called_once() + mock_prepare_bindflow.assert_not_called() + assert ( + mock_prepare_proteindj.call_args.kwargs["pipeline"] == "https://github.com/test/proteindj" + ) + assert mock_prepare_proteindj.call_args.kwargs["output_id"] == data["runId"] + + def test_launch_invalid_payload(client: TestClient): """Test launch with invalid payload.""" payload = { @@ -477,59 +569,6 @@ def test_extract_final_design_count_string_number(): assert _extract_final_design_count(_form_data(number_of_final_designs="25")) == 25 -# ============================================================================= -# Tests for missing repo_url / default_revision -# ============================================================================= - - -def test_launch_missing_repo_url(client: TestClient, app, test_engine): - """Workflow missing repo_url should return 500.""" - with Session(test_engine) as db: - db.add( - Workflow( - id=uuid4(), - name="single-prediction", - description="No repo workflow", - repo_url=None, - default_revision="dev", - ) - ) - db.commit() - - payload = { - "launch": {"workflow": "single-prediction", "tool": "colabfold", "runName": "test-run"}, - "s3InputKey": "inputs/samplesheets/test.csv", - "formData": {"workflow": "single-prediction", "tool": "colabfold"}, - } - response = client.post("/api/workflows/launch", json=payload) - assert response.status_code == 500 - assert "missing repo_url" in response.json()["detail"] - - -def test_launch_missing_default_revision(client: TestClient, app, test_engine): - """Workflow missing default_revision should return 500.""" - with Session(test_engine) as db: - db.add( - Workflow( - id=uuid4(), - name="single-prediction", - description="No revision workflow", - repo_url="https://github.com/test/norev", - default_revision=None, - ) - ) - db.commit() - - payload = { - "launch": {"workflow": "single-prediction", "tool": "colabfold", "runName": "test-run"}, - "s3InputKey": "inputs/samplesheets/test.csv", - "formData": {"workflow": "single-prediction", "tool": "colabfold"}, - } - response = client.post("/api/workflows/launch", json=payload) - assert response.status_code == 500 - assert "missing default_revision" in response.json()["detail"] - - # ============================================================================= # Tests for proteinfold launch path # ============================================================================= @@ -564,7 +603,18 @@ def test_launch_proteinfold_success(mock_prepare, client: TestClient, test_engin payload = { "launch": {"workflow": "single-prediction", "tool": "colabfold", "runName": "pf-run-1"}, "s3InputKey": "inputs/samplesheets/test.csv", - "formData": {"workflow": "single-prediction", "tool": "colabfold"}, + "formData": { + "workflow": "single-prediction", + "tool": "colabfold", + "entities": [ + { + "id": "seq1", + "moleculeType": "protein", + "copyNumber": 1, + "sequence": "ACDEFGHIK", + } + ], + }, } response = client.post("/api/workflows/launch", json=payload) @@ -600,7 +650,18 @@ def test_launch_proteinfold_queue_preparation_configuration_error( "runName": "pf-run-cfg-err", }, "s3InputKey": "inputs/samplesheets/test.csv", - "formData": {"workflow": "single-prediction", "tool": "colabfold"}, + "formData": { + "workflow": "single-prediction", + "tool": "colabfold", + "entities": [ + { + "id": "seq1", + "moleculeType": "protein", + "copyNumber": 1, + "sequence": "ACDEFGHIK", + } + ], + }, } response = client.post("/api/workflows/launch", json=payload) @@ -621,7 +682,18 @@ def test_launch_proteinfold_queue_preparation_error(mock_prepare, client: TestCl "runName": "pf-run-exec-err", }, "s3InputKey": "inputs/samplesheets/test.csv", - "formData": {"workflow": "single-prediction", "tool": "colabfold"}, + "formData": { + "workflow": "single-prediction", + "tool": "colabfold", + "entities": [ + { + "id": "seq1", + "moleculeType": "protein", + "copyNumber": 1, + "sequence": "ACDEFGHIK", + } + ], + }, } response = client.post("/api/workflows/launch", json=payload) @@ -629,6 +701,91 @@ def test_launch_proteinfold_queue_preparation_error(mock_prepare, client: TestCl assert response.json()["detail"] == "Failed to queue local workflow run." +# ============================================================================= +# Tests for single-prediction entity validation +# ============================================================================= + + +def _single_prediction_payload(entities, tool="colabfold", **form_extra): + form_data = {"workflow": "single-prediction", "tool": tool, "entities": entities} + form_data.update(form_extra) + return { + "launch": {"workflow": "single-prediction", "tool": tool, "runName": "sp-val"}, + "s3InputKey": "inputs/samplesheets/test.csv", + "formData": form_data, + } + + +def _protein_entity(sequence="ACDEFGHIK", copy_number=1): + return { + "id": "seq1", + "moleculeType": "protein", + "copyNumber": copy_number, + "sequence": sequence, + } + + +def test_launch_single_prediction_rejects_missing_entities(client: TestClient, test_engine): + _add_proteinfold_workflow(test_engine) + payload = { + "launch": {"workflow": "single-prediction", "tool": "colabfold", "runName": "sp-no-ent"}, + "s3InputKey": "inputs/samplesheets/test.csv", + "formData": {"workflow": "single-prediction", "tool": "colabfold"}, + } + response = client.post("/api/workflows/launch", json=payload) + assert response.status_code == 422 + assert "entities" in response.json()["detail"] + + +def test_launch_single_prediction_rejects_too_many_entities(client: TestClient, test_engine): + _add_proteinfold_workflow(test_engine) + payload = _single_prediction_payload([_protein_entity(copy_number=53)]) + response = client.post("/api/workflows/launch", json=payload) + assert response.status_code == 422 + assert "Too many entities" in response.json()["detail"] + + +def test_launch_single_prediction_requires_protein(client: TestClient, test_engine): + _add_proteinfold_workflow(test_engine) + entities = [ + {"id": "seq1", "moleculeType": "dna", "copyNumber": 1, "sequence": "ACGT"}, + ] + payload = _single_prediction_payload(entities, tool="boltz") + response = client.post("/api/workflows/launch", json=payload) + assert response.status_code == 422 + assert "must be a protein" in response.json()["detail"] + + +def test_launch_single_prediction_rejects_oversized_alphafold2(client: TestClient, test_engine): + _add_proteinfold_workflow(test_engine) + payload = _single_prediction_payload([_protein_entity(sequence="A" * 2000)], tool="alphafold2") + response = client.post("/api/workflows/launch", json=payload) + assert response.status_code == 422 + assert "less than 2000" in response.json()["detail"] + + +@patch( + "app.routes.workflows.prepare_proteinfold_workflow", side_effect=_queue_job_for_route_prepare +) +def test_launch_single_prediction_boltz_potentials_reduces_limit( + mock_prepare, client: TestClient, test_engine +): + _add_proteinfold_workflow(test_engine) + + ok_payload = _single_prediction_payload( + [_protein_entity(sequence="A" * 1999)], tool="boltz", boltz_use_potentials=True + ) + assert client.post("/api/workflows/launch", json=ok_payload).status_code == 201 + + over_payload = _single_prediction_payload( + [_protein_entity(sequence="A" * 2000)], tool="boltz", boltz_use_potentials=True + ) + response = client.post("/api/workflows/launch", json=over_payload) + assert response.status_code == 422 + assert "less than 2000" in response.json()["detail"] + mock_prepare.assert_called_once() + + # ============================================================================= # Tests for require_workflow_execution_role # ============================================================================= @@ -771,62 +928,6 @@ def _get_db(): yield c -@pytest.fixture -def wisps_no_config_client(test_engine): - """Test client with an interaction-screening workflow that has config_path=None.""" - from sqlalchemy.orm import sessionmaker - - from app.main import create_app - - application = create_app() - user_id = UUID("11111111-1111-1111-1111-111111111111") - - SessionLocal = sessionmaker( - bind=test_engine, autocommit=False, autoflush=False, expire_on_commit=False - ) - setup_session = SessionLocal() - - if not setup_session.get(AppUser, user_id): - setup_session.add( - AppUser( - id=user_id, - auth0_user_id="auth0|test-user", - name="Test User", - email="test@example.com", - ) - ) - - setup_session.add( - Workflow( - id=uuid4(), - name="interaction-screening", - description="WISPS workflow without config_path", - repo_url="https://github.com/test/wisps", - default_revision="main", - config_path=None, - ) - ) - - setup_session.commit() - setup_session.close() - - def _get_db(): - db = SessionLocal() - try: - yield db - finally: - db.close() - - from app.routes.dependencies import require_workflow_execution_role - - application.dependency_overrides[get_db] = _get_db - application.dependency_overrides[get_current_user_id] = lambda: user_id - application.dependency_overrides[require_workflow_execution_role] = lambda: None - - with TestClient(application) as c: - yield c - - @patch("app.routes.workflows.prepare_wisps_workflow", side_effect=_queue_job_for_route_prepare) def test_launch_interaction_screening_success(mock_prepare, wisps_client: TestClient, test_engine): """Test successful interaction-screening workflow launch.""" @@ -994,29 +1095,6 @@ def test_launch_interaction_screening_queue_preparation_error( assert count == 0 -def test_launch_interaction_screening_missing_config_path(wisps_no_config_client: TestClient): - """interaction-screening workflow missing config_path should return 500.""" - payload = { - "launch": { - "workflow": "interaction-screening", - "tool": "boltz", - "runName": "wisps-run-no-cfg", - }, - "s3InputKey": "inputs/samplesheets/test.csv", - "formData": { - "workflow": "interaction-screening", - "tool": "boltz", - "fastaS3Uri": "s3://bucket/test.fasta", - "splitOutputDir": "/data/split", - }, - } - - response = wisps_no_config_client.post("/api/workflows/launch", json=payload) - - assert response.status_code == 500 - assert "config_path" in response.json()["detail"] - - @patch("app.routes.workflows.prepare_wisps_workflow", side_effect=_queue_job_for_route_prepare) def test_launch_with_workflow_field_in_launch(mock_prepare, wisps_client: TestClient, test_engine): """The new frontend format using launch.workflow is accepted alongside launch.tool.""" diff --git a/tests/test_schemas.py b/tests/test_schemas.py index fbe5200..a7a5a80 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -17,13 +17,17 @@ ListRunsResponse, PipelineStatus, RunInfo, + SinglePredictionEntity, UIStatus, WispsDatasetUploadRequest, WispsSequenceItem, WorkflowLaunchForm, WorkflowLaunchPayload, WorkflowLaunchResponse, + WorkflowUserDetails, map_pipeline_status_to_ui, + single_prediction_size_limit, + validate_single_prediction_entities, ) @@ -462,3 +466,89 @@ def test_interaction_screening_request_extra_fields_forbidden(): runId="run-1", extra="bad", ) + + +def test_get_encoded_account_details_encodes_email_and_ip(): + details = WorkflowUserDetails(user_email="user@example.com", ip_address="1.2.3.4") + assert details.get_encoded_account_details() == "dXNlckBleGFtcGxlLmNvbQ==:MS4yLjMuNA==" + + +def test_get_encoded_account_details_without_ip_encodes_email_only(): + details = WorkflowUserDetails(user_email="user@example.com", ip_address="") + assert details.get_encoded_account_details() == "dXNlckBleGFtcGxlLmNvbQ==" + + +# ============================================================================= +# Single-prediction entity validation +# ============================================================================= + + +def _protein(sequence="ACDEFGHIK", copy_number=1): + return SinglePredictionEntity( + id="seq1", moleculeType="protein", copyNumber=copy_number, sequence=sequence + ) + + +def test_size_limit_defaults_and_boltz_potentials(): + assert single_prediction_size_limit("colabfold", False) == 4000 + assert single_prediction_size_limit("boltz", False) == 4000 + assert single_prediction_size_limit("alphafold2", False) == 2000 + assert single_prediction_size_limit("boltz", True) == 2000 + assert single_prediction_size_limit("colabfold", True) == 4000 + + +def test_validate_single_prediction_accepts_valid_input(): + validate_single_prediction_entities([_protein()], "colabfold") + + +def test_validate_single_prediction_requires_entities(): + with pytest.raises(ValueError, match="At least one entity"): + validate_single_prediction_entities([], "colabfold") + + +def test_validate_single_prediction_counts_copies_toward_limit(): + validate_single_prediction_entities([_protein(copy_number=52)], "colabfold") + with pytest.raises(ValueError, match="Too many entities"): + validate_single_prediction_entities([_protein(copy_number=53)], "colabfold") + + +def test_validate_single_prediction_requires_protein(): + dna = SinglePredictionEntity(id="d", moleculeType="dna", copyNumber=1, sequence="ACGT") + with pytest.raises(ValueError, match="must be a protein"): + validate_single_prediction_entities([dna], "boltz") + + +def test_single_prediction_entity_rejects_copy_number_below_one(): + with pytest.raises(ValidationError): + SinglePredictionEntity(id="p", moleculeType="protein", copyNumber=0, sequence="ACD") + + +def test_single_prediction_entity_rejects_unknown_molecule_type(): + with pytest.raises(ValidationError): + SinglePredictionEntity(id="p", moleculeType="peptide", copyNumber=1, sequence="ACD") + + +def test_single_prediction_entity_rejects_empty_id(): + with pytest.raises(ValidationError): + SinglePredictionEntity(id="", moleculeType="protein", copyNumber=1, sequence="ACD") + + +def test_validate_single_prediction_ligand_uses_fixed_size(): + protein = _protein(sequence="A" * 3960) + ligand = SinglePredictionEntity( + id="lig", moleculeType="ligand", copyNumber=1, sequence="CC(=O)O" + ) + # 3960 + 30 (ligand) = 3990 < 4000 + validate_single_prediction_entities([protein, ligand], "colabfold") + # Two ligand copies: 3960 + 60 = 4020 >= 4000 + ligand_two = SinglePredictionEntity( + id="lig", moleculeType="ligand", copyNumber=2, sequence="CC(=O)O" + ) + with pytest.raises(ValueError, match="less than 4000"): + validate_single_prediction_entities([protein, ligand_two], "colabfold") + + +def test_validate_single_prediction_size_limit_is_exclusive(): + with pytest.raises(ValueError, match="less than 2000"): + validate_single_prediction_entities([_protein(sequence="A" * 2000)], "alphafold2") + validate_single_prediction_entities([_protein(sequence="A" * 1999)], "alphafold2") diff --git a/tests/test_services_bindflow_config.py b/tests/test_services_bindflow_config.py index 3ff9f77..f1f69d3 100644 --- a/tests/test_services_bindflow_config.py +++ b/tests/test_services_bindflow_config.py @@ -18,12 +18,11 @@ from app.services.launch_payloads import get_executor_script _SHEET_URL = "https://api.seqera.test/workspaces/ws1/datasets/ds1/v/1/n/samplesheet.csv" -_USER_DETAILS = WorkflowUserDetails( - user_email="user@example.com", - full_name="", - institute="", - ip_address="", -) + + +def _user_details(email: str, ip_address: str = "") -> WorkflowUserDetails: + return WorkflowUserDetails(user_email=email, ip_address=ip_address) + # ============================================================================= # Tests for get_bindflow_default_params() @@ -125,9 +124,7 @@ def test_get_bindflow_config_text_includes_base_config(): with patch("builtins.open", mock_open(read_data="base_config_content")): result = get_bindflow_config_text( "/fake/bindflow.config", - job_id="run-1", - user_details=_USER_DETAILS, - timestamp="20260507_120000", + user_details=_user_details("user@example.com"), ) assert "base_config_content" in result @@ -136,43 +133,38 @@ def test_get_bindflow_config_text_appends_process_block(): with patch("builtins.open", mock_open(read_data="")): result = get_bindflow_config_text( "/fake/bindflow.config", - job_id="run-1", - user_details=_USER_DETAILS, - timestamp="20260507_120000", + user_details=_user_details("user@example.com"), ) assert "process {" in result assert "clusterOptions" in result -def test_get_bindflow_config_text_interpolates_job_fields(): +def test_get_bindflow_config_text_interpolates_email(): + with patch("builtins.open", mock_open(read_data="")): + result = get_bindflow_config_text( + "/fake/bindflow.config", + user_details=_user_details("alice@example.com"), + ) + assert "-A YWxpY2VAZXhhbXBsZS5jb20=" in result + + +def test_get_bindflow_config_text_without_ip_address_omits_encoding(): with patch("builtins.open", mock_open(read_data="")): result = get_bindflow_config_text( "/fake/bindflow.config", - job_id="my-run", - user_details=_USER_DETAILS.model_copy(update={"user_email": "alice@example.com"}), - timestamp="20260507_090000", + user_details=_user_details("user@example.com"), ) - assert "my-run" in result - assert "alice@example.com" in result - assert "20260507_090000" in result + assert "-A dXNlckBleGFtcGxlLmNvbQ==" in result + assert ":" not in result.split("clusterOptions = ")[1] -def test_get_bindflow_config_text_optional_fields(): +def test_get_bindflow_config_text_with_ip_address_appends_encoded_ip(): with patch("builtins.open", mock_open(read_data="")): result = get_bindflow_config_text( "/fake/bindflow.config", - job_id="run-1", - user_details=WorkflowUserDetails( - user_email="user@example.com", - full_name="Alice Smith", - institute="BioCommons", - ip_address="1.2.3.4", - ), - timestamp="ts", + user_details=_user_details("user@example.com", ip_address="1.2.3.4"), ) - assert "Alice Smith" in result - assert "BioCommons" in result - assert "1.2.3.4" in result + assert "-A dXNlckBleGFtcGxlLmNvbQ==:MS4yLjMuNA==" in result def test_get_bindflow_config_text_url_fetching(): @@ -182,9 +174,7 @@ def test_get_bindflow_config_text_url_fetching(): ) result = get_bindflow_config_text( "https://raw.githubusercontent.com/org/repo/main/bindflow.config", - job_id="run-url", - user_details=_USER_DETAILS, - timestamp="20260507_120000", + user_details=_user_details("user@example.com"), ) assert "remote_base_config" in result assert "clusterOptions" in result @@ -198,7 +188,5 @@ def test_get_bindflow_config_text_url_error_raises(): with pytest.raises(httpx.HTTPStatusError): get_bindflow_config_text( "https://raw.githubusercontent.com/org/repo/main/bindflow.config", - job_id="run-1", - user_details=_USER_DETAILS, - timestamp="ts", + user_details=_user_details("user@example.com"), ) diff --git a/tests/test_services_cluster_utils.py b/tests/test_services_cluster_utils.py new file mode 100644 index 0000000..f4fc80c --- /dev/null +++ b/tests/test_services_cluster_utils.py @@ -0,0 +1,27 @@ +"""Tests for cluster_utils.py.""" + +# pylint: disable=missing-function-docstring +from __future__ import annotations + +import importlib + +from app.services import cluster_utils + + +def test_gadi_project_defaults_to_yz52(monkeypatch): + monkeypatch.delenv("GADI_PROJECT", raising=False) + reloaded = importlib.reload(cluster_utils) + try: + assert reloaded.GADI_PROJECT == "yz52" + finally: + importlib.reload(cluster_utils) + + +def test_gadi_project_reads_env_var(monkeypatch): + monkeypatch.setenv("GADI_PROJECT", "ab12") + reloaded = importlib.reload(cluster_utils) + try: + assert reloaded.GADI_PROJECT == "ab12" + finally: + monkeypatch.delenv("GADI_PROJECT", raising=False) + importlib.reload(cluster_utils) diff --git a/tests/test_services_health.py b/tests/test_services_health.py index 47a6d20..cff67e4 100644 --- a/tests/test_services_health.py +++ b/tests/test_services_health.py @@ -5,8 +5,10 @@ from datetime import UTC, datetime, timedelta import httpx +import pytest +from sqlalchemy import select -from app.db.models.system_status import SystemStatusCache +from app.db.models.system_status import SystemStatusCache, SystemStatusIncident from app.schemas.health import ( COMPONENT_COMPUTE_ENV, COMPONENT_SEQERA_API, @@ -15,6 +17,19 @@ from app.services import health +@pytest.fixture(autouse=True) +def _agent_healthcheck_disabled_by_default(monkeypatch): + """Keep these tests hermetic regardless of a developer's local .env. + + The Tower Agent probe is opt-in via ENABLE_AGENT_HEALTHCHECK, and app.main + loads .env at import time, so a developer with it set locally would + otherwise get an unmocked probe_tower_agent() running in every test here. + Tests that specifically exercise the agent probe re-enable it themselves + (see _install_agent_mock). + """ + monkeypatch.delenv("ENABLE_AGENT_HEALTHCHECK", raising=False) + + def _component(status: health.SystemStatus, name: str) -> health.ProbeResult: return next(c for c in status.components if c.name == name) @@ -216,6 +231,201 @@ async def ok_get(self, url, *args, **kwargs): # noqa: ANN001 assert stale.checked_at == first.checked_at +# --------------------------------------------------------------------------- +# Downtime incidents (only degraded/unhealthy periods are recorded) +# --------------------------------------------------------------------------- + + +def _incidents(db, component=None): + rows = db.execute(select(SystemStatusIncident)).scalars().all() + return [r for r in rows if component is None or r.component == component] + + +async def test_healthy_refresh_writes_no_incidents(monkeypatch, test_db): + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("AVAILABLE"), + ) + await health.get_system_status(test_db, force_refresh=True) + + assert _incidents(test_db) == [] + + +async def test_going_unhealthy_opens_an_incident(monkeypatch, test_db): + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("ERRORED", "Gadi agent disconnected"), + ) + status = await health.get_system_status(test_db, force_refresh=True) + + rows = _incidents(test_db, COMPONENT_COMPUTE_ENV) + assert len(rows) == 1 + assert rows[0].status == "unhealthy" + assert rows[0].ended_at is None + assert "Gadi agent disconnected" in (rows[0].message or "") + assert rows[0].started_at == status.checked_at.replace(tzinfo=None) + + +async def test_recovering_to_healthy_closes_the_incident(monkeypatch, test_db): + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("ERRORED"), + ) + await health.get_system_status(test_db, force_refresh=True) + + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("AVAILABLE"), + ) + recovered = await health.get_system_status(test_db, force_refresh=True) + + rows = _incidents(test_db, COMPONENT_COMPUTE_ENV) + assert len(rows) == 1 # closed, not duplicated + assert rows[0].ended_at == recovered.checked_at.replace(tzinfo=None) + + +async def test_staying_unhealthy_does_not_open_a_second_incident(monkeypatch, test_db): + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("ERRORED", "first error"), + ) + await health.get_system_status(test_db, force_refresh=True) + + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("ERRORED", "still erroring"), + ) + await health.get_system_status(test_db, force_refresh=True) + + rows = _incidents(test_db, COMPONENT_COMPUTE_ENV) + assert len(rows) == 1 + assert rows[0].ended_at is None + # The message on the still-open incident is refreshed to the latest probe result. + assert rows[0].message == "Compute environment state: ERRORED (still erroring)" + + +async def test_severity_change_splits_the_incident(monkeypatch, test_db): + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("CREATING"), # degraded + ) + await health.get_system_status(test_db, force_refresh=True) + + _mock_response( + monkeypatch, + user_info=_ok_user_info, + compute_env=_compute_env_with_status("ERRORED"), # unhealthy + ) + await health.get_system_status(test_db, force_refresh=True) + + rows = sorted(_incidents(test_db, COMPONENT_COMPUTE_ENV), key=lambda r: r.started_at) + assert len(rows) == 2 + assert rows[0].status == "degraded" + assert rows[0].ended_at is not None # closed when severity changed + assert rows[1].status == "unhealthy" + assert rows[1].ended_at is None + + +def test_get_incidents_filters_by_window_and_component(test_db): + now = datetime.now(UTC) + long_ago = SystemStatusIncident( + component=COMPONENT_SEQERA_API, + status="unhealthy", + started_at=now - timedelta(hours=5), + ended_at=now - timedelta(hours=4, minutes=55), + ) + recent_api = SystemStatusIncident( + component=COMPONENT_SEQERA_API, + status="unhealthy", + started_at=now - timedelta(minutes=30), + ended_at=now - timedelta(minutes=25), + ) + ongoing_ce = SystemStatusIncident( + component=COMPONENT_COMPUTE_ENV, + status="degraded", + started_at=now - timedelta(minutes=10), + ended_at=None, + ) + test_db.add_all([long_ago, recent_api, ongoing_ce]) + test_db.commit() + + since = now - timedelta(hours=1) + all_recent = health.get_incidents(test_db, since=since, until=now) + assert {row.component for row in all_recent} == {COMPONENT_SEQERA_API, COMPONENT_COMPUTE_ENV} + + api_only = health.get_incidents(test_db, since=since, until=now, component=COMPONENT_SEQERA_API) + assert len(api_only) == 1 + assert api_only[0].ended_at is not None + + +def test_summarize_downtime_clips_to_window(): + now = datetime.now(UTC) + since = now - timedelta(hours=1) + incidents = [ + # Started before the window, ended inside it -> clipped at `since`. + SystemStatusIncident( + component="x", + status="unhealthy", + started_at=now - timedelta(hours=2), + ended_at=since + timedelta(minutes=10), + ), + # Still ongoing -> counts through `until`. + SystemStatusIncident( + component="x", + status="degraded", + started_at=now - timedelta(minutes=5), + ended_at=None, + ), + ] + summary = health._summarize_downtime(incidents, since=since, until=now) + + assert summary["incidentCount"] == 2 + assert summary["downtimeSeconds"] == 10 * 60 + 5 * 60 + assert summary["uptimePercent"] == round((3600 - 900) / 3600 * 100, 2) + + +def test_summarize_downtime_no_incidents_is_full_uptime(): + now = datetime.now(UTC) + summary = health._summarize_downtime([], since=now - timedelta(hours=1), until=now) + assert summary == {"incidentCount": 0, "downtimeSeconds": 0.0, "uptimePercent": 100.0} + + +def test_to_downtime_dict_groups_by_component(): + now = datetime.now(UTC) + since = now - timedelta(hours=24) + incidents = [ + SystemStatusIncident( + component=COMPONENT_SEQERA_API, + status="unhealthy", + started_at=since + timedelta(hours=1), + ended_at=since + timedelta(hours=2), + ), + SystemStatusIncident( + component=COMPONENT_COMPUTE_ENV, + status="degraded", + started_at=since + timedelta(hours=3), + ended_at=None, + message="still starting up", + ), + ] + result = health.to_downtime_dict(incidents, since=since, until=now, hours=24) + + assert result["windowHours"] == 24 + assert result["since"] == since + assert result["until"] == now + by_name = {c["name"]: c for c in result["components"]} + assert by_name[COMPONENT_SEQERA_API]["summary"]["incidentCount"] == 1 + assert by_name[COMPONENT_SEQERA_API]["incidents"][0]["endedAt"] == since + timedelta(hours=2) + assert by_name[COMPONENT_COMPUTE_ENV]["incidents"][0]["message"] == "still starting up" + + def test_overall_status_aggregation(): assert health._worst(["healthy", "healthy"]) == "healthy" assert health._worst(["healthy", "degraded"]) == "degraded" @@ -321,7 +531,6 @@ async def fake_delete(self, url, *args, **kwargs): # noqa: ANN001 async def test_agent_probe_disabled_by_default(monkeypatch): - monkeypatch.delenv("ENABLE_AGENT_HEALTHCHECK", raising=False) _mock_response( monkeypatch, user_info=_ok_user_info, diff --git a/tests/test_services_job_utils.py b/tests/test_services_job_utils.py index 4fe5ad6..c4a169d 100644 --- a/tests/test_services_job_utils.py +++ b/tests/test_services_job_utils.py @@ -43,7 +43,12 @@ def commit(self): def _configure_bindcraft_run(run: WorkflowRun) -> None: - run.workflow = Workflow(name="de-novo-design") + run.workflow = Workflow( + name="de-novo-design", + repo_url="https://github.com/test/de-novo-design", + default_revision="main", + config_path="/config/de-novo-design.config", + ) run.tool = "bindcraft" run.submitted_form_data = {"mode": "bindcraft"} @@ -69,8 +74,20 @@ def test_get_user_job_list_rows_returns_current_user_run_metadata(test_db): name="Other User", email="other-jobs@example.com", ) - workflow = Workflow(name="de-novo-design", description="Design workflow") - other_workflow = Workflow(name="other-workflow", description="Other workflow") + workflow = Workflow( + name="de-novo-design", + description="Design workflow", + repo_url="https://github.com/test/de-novo-design", + default_revision="main", + config_path="/config/de-novo-design.config", + ) + other_workflow = Workflow( + name="other-workflow", + description="Other workflow", + repo_url="https://github.com/test/other-workflow", + default_revision="main", + config_path="/config/other-workflow.config", + ) test_db.add_all([user, other_user, workflow, other_workflow]) test_db.commit() @@ -243,7 +260,12 @@ async def test_ensure_completed_run_score_persists_spec_score(test_db): name="Score User", email="score-user@example.com", ) - workflow = Workflow(name="de-novo-design") + workflow = Workflow( + name="de-novo-design", + repo_url="https://github.com/test/de-novo-design", + default_revision="main", + config_path="/config/de-novo-design.config", + ) run = WorkflowRun( owner=user, workflow=workflow, @@ -472,7 +494,12 @@ async def test_get_result_snapshot_downloads_skips_s3_for_proteinfold_workflows( name="Proteinfold No Snapshot User", email="proteinfold-no-snapshot-user@example.com", ) - workflow = Workflow(name="single-prediction") + workflow = Workflow( + name="single-prediction", + repo_url="https://github.com/test/single-prediction", + default_revision="main", + config_path="/config/single-prediction.config", + ) run = WorkflowRun( owner=user, workflow=workflow, diff --git a/tests/test_services_results_utils.py b/tests/test_services_results_utils.py index de2732b..9db273d 100644 --- a/tests/test_services_results_utils.py +++ b/tests/test_services_results_utils.py @@ -4,7 +4,7 @@ from io import BytesIO from types import SimpleNamespace -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 from zipfile import ZipFile @@ -20,29 +20,36 @@ build_bindcraft_output_listing_prefixes, build_boltz_proteinfold_output_listing_prefixes, build_colabfold_proteinfold_output_listing_prefixes, + build_rfdiffusion_output_listing_prefixes, build_wisps_output_listing_prefixes, classify_alphafold2_proteinfold_output, classify_bindcraft_output_key, classify_boltz_proteinfold_output, classify_colabfold_proteinfold_output, + classify_rfdiffusion_output_key, + classify_shared_outputs, classify_wisps_output_key, extract_bindcraft_max_score, extract_proteinfold_max_score, + extract_rfdiffusion_max_score, extract_wisps_max_score, format_log_entries, get_all_downloads_zipped, get_bindcraft_score_file, get_proteinfold_score_file, + get_rfdiffusion_score_file, + get_run_service_usage, get_sample_id_for_result, get_tool_name, get_wisps_score_file, + list_workflow_outputs_from_s3, resolve_fasta_form_data, resolve_pdb_presigned_urls, resolve_submitted_form_data, s3_uri_to_key, ) from app.services.s3 import S3ServiceError -from tests.datagen import AppUserFactory, WorkflowRunFactory +from tests.datagen import AppUserFactory, RunOutputFactory, S3ObjectFactory, WorkflowRunFactory def test_format_log_entries_extracts_timestamp_level_and_strips_ansi(): @@ -210,6 +217,7 @@ def test_bindcraft_helpers_classify_keys_and_build_prefixes(monkeypatch): prefixes = build_bindcraft_output_listing_prefixes(run) assert prefixes == [ + f"{run.id}/", f"{run.id}/ranker/", f"{run.id}/generate/", f"{run.id}/bindcraft/sampleZ_0_output/", @@ -217,6 +225,7 @@ def test_bindcraft_helpers_classify_keys_and_build_prefixes(monkeypatch): run_without_sample = SimpleNamespace(id=run.id, sample_id=None, binder_name=None, form_id=None) assert build_bindcraft_output_listing_prefixes(run_without_sample) == [ + f"{run.id}/", f"{run.id}/ranker/", f"{run.id}/generate/", ] @@ -227,6 +236,81 @@ def test_bindcraft_helpers_classify_keys_and_build_prefixes(monkeypatch): assert _build_s3_uri("path/to/file.txt") == "path/to/file.txt" +def test_rfdiffusion_helpers_classify_keys_and_build_prefixes(): + run = WorkflowRun(id=uuid4(), owner_user_id=uuid4(), sample_id="sampleZ") + + assert classify_rfdiffusion_output_key(" ") is None + assert classify_rfdiffusion_output_key("folder/") is None + assert classify_rfdiffusion_output_key(f"{run.id}/ranked_designs.csv") is None + assert classify_rfdiffusion_output_key(f"{run.id}/results/other.txt") is None + assert classify_rfdiffusion_output_key( + f"{run.id}/results/ranked_designs.csv" + ) == ClassifiedOutput( + "stats_csv", + "ranked_designs.csv", + ) + assert classify_rfdiffusion_output_key( + f"{run.id}/results/ranked_designs.tar.gz" + ) == ClassifiedOutput( + "pdb", + "ranked_designs.tar.gz", + ) + + assert build_rfdiffusion_output_listing_prefixes(run) == [ + f"{run.id}/", + f"{run.id}/results/", + ] + assert build_rfdiffusion_output_listing_prefixes(SimpleNamespace(id=None)) == [] + + +def test_get_rfdiffusion_score_file_uses_ranked_designs_csv(): + assert ( + get_rfdiffusion_score_file( + [ + " ", + "run-1/results/other.csv", + " run-1/results/ranked_designs.csv ", + ], + sample_id="sampleZ", + ) + == "run-1/results/ranked_designs.csv" + ) + assert get_rfdiffusion_score_file(["run-1/results/other.csv"], sample_id=None) is None + + +@pytest.mark.asyncio +async def test_extract_rfdiffusion_max_score_reads_first_ranked_design_score(): + csv_text = "design,af2_plddt_overall\nsampleZ_0,91.3\nsampleZ_1,88.1\n" + + with patch( + "app.services.results_utils.read_s3_file", + new_callable=AsyncMock, + return_value=csv_text, + ) as read_file: + score = await extract_rfdiffusion_max_score("run-1/results/ranked_designs.csv") + + assert score == pytest.approx(0.913) + read_file.assert_awaited_once_with("run-1/results/ranked_designs.csv") + + +@pytest.mark.asyncio +async def test_extract_rfdiffusion_max_score_returns_none_without_score_value(): + """ + Test extract_rfdiffusion_max_score returns None when no af2_plddt_overall score is available + """ + for csv_text in [ + "design,af2_plddt_overall\n", + "design,af2_plddt_overall\nsampleZ_0, \n", + "design,other_score\nsampleZ_0,91.3\n", + ]: + with patch( + "app.services.results_utils.read_s3_file", + new_callable=AsyncMock, + return_value=csv_text, + ): + assert await extract_rfdiffusion_max_score("run-1/results/ranked_designs.csv") is None + + @pytest.mark.asyncio async def test_get_all_downloads_zipped_writes_category_label_files_and_reads_each_output( test_db, persistent_models @@ -234,7 +318,12 @@ async def test_get_all_downloads_zipped_writes_category_label_files_and_reads_ea user = AppUserFactory.create_sync() run = WorkflowRunFactory.create_sync( owner=user, - workflow=Workflow(name="de-novo-design"), + workflow=Workflow( + name="de-novo-design", + repo_url="https://github.com/test/de-novo-design", + default_revision="main", + config_path="/config/de-novo-design.config", + ), tool="bindcraft", seqera_run_id="wf-zip-results", ) @@ -347,6 +436,51 @@ def test_all_workflow_output_specs_have_score_hooks(): assert spec.extract_max_score is not None +@pytest.mark.asyncio +async def test_list_workflow_outputs_from_s3_continues(): + """ + Test list_wokflow_outputs_from_s3 continues past S3 errors when suppress_s3_errors is True + """ + run = WorkflowRun( + id=uuid4(), + owner_user_id=uuid4(), + sample_id="sampleZ", + seqera_run_id="seqera-run-1", + ) + + async def extract_max_score(_key: str) -> float | None: + return None + + spec = WorkflowResultsSpec( + kind="de-novo-design", + tool="rfdiffusion", + required_categories=set(), + get_prefixes=lambda _run: ["failed-prefix/", "ok-prefix/"], + classifier=classify_rfdiffusion_output_key, + get_score_file=lambda _keys, _sample_id: None, + extract_max_score=extract_max_score, + ) + + with patch( + "app.services.results_utils.list_s3_files", + new=AsyncMock( + side_effect=[ + S3ServiceError("failed to list"), + [{"key": "ok-prefix/results/ranked_designs.csv"}], + ] + ), + ) as list_s3_files_mock: + outputs = await list_workflow_outputs_from_s3(run, spec, suppress_s3_errors=True) + + assert outputs == { + "ok-prefix/results/ranked_designs.csv": ClassifiedOutput( + category="stats_csv", + label="ranked_designs.csv", + ) + } + assert list_s3_files_mock.await_count == 2 + + @pytest.mark.asyncio async def test_workflow_results_spec_get_max_score_returns_none_without_score_file(test_db): user = AppUser( @@ -369,7 +503,7 @@ async def test_workflow_results_spec_get_max_score_returns_none_without_score_fi tool="boltz", required_categories=set(), get_prefixes=lambda _run: [], - classify=lambda _key, _sample_id: None, + classifier=lambda _key, _sample_id: None, get_score_file=lambda _keys, _sample_id: None, extract_max_score=extractor, ) @@ -415,7 +549,7 @@ async def test_workflow_results_spec_get_max_score_extracts_selected_run_output( tool="boltz", required_categories=set(), get_prefixes=lambda _run: [], - classify=lambda _key, _sample_id: None, + classifier=lambda _key, _sample_id: None, get_score_file=get_proteinfold_score_file, extract_max_score=extractor, ) @@ -424,6 +558,50 @@ async def test_workflow_results_spec_get_max_score_extracts_selected_run_output( extractor.assert_awaited_once_with("run-1/boltz/T1024/T1024_ptm.tsv") +@pytest.mark.asyncio +async def test_workflow_results_spec_get_service_units_uses_usage_report_key( + test_db, persistent_models +): + user = AppUserFactory.create_sync() + run = WorkflowRunFactory.create_sync( + owner=user, + work_dir="workdir-usage-selected", + ) + ignored = S3ObjectFactory.create_sync( + object_key="run-1/boltz/T1024/T1024_ptm.tsv", + uri="s3://bucket/run-1/boltz/T1024/T1024_ptm.tsv", + ) + usage_object = S3ObjectFactory.create_sync( + object_key="run-1/UsageReport.csv", + uri="s3://bucket/run-1/UsageReport.csv", + ) + RunOutputFactory.create_sync(run_id=run.id, s3_object_id=ignored.object_key) + RunOutputFactory.create_sync( + run_id=run.id, + s3_object_id=usage_object.object_key, + ) + + spec = WorkflowResultsSpec( + kind="single-prediction", + tool="boltz", + required_categories=set(), + get_prefixes=lambda _run: [], + classifier=lambda _key, _sample_id: None, + get_score_file=get_proteinfold_score_file, + extract_max_score=AsyncMock(return_value=0.91), + ) + + with patch( + "app.services.results_utils.get_run_service_usage", + new_callable=AsyncMock, + return_value=3.45, + ) as get_run_service_usage_mock: + service_units = await spec.get_service_units(test_db, run) + + assert service_units == 3.45 + get_run_service_usage_mock.assert_awaited_once_with("run-1/UsageReport.csv") + + def test_boltz_proteinfold_helpers_classify_keys_and_build_prefixes(): run = WorkflowRun(id=uuid4(), owner_user_id=uuid4(), sample_id="T1024") @@ -437,6 +615,9 @@ def test_boltz_proteinfold_helpers_classify_keys_and_build_prefixes(): assert classify_boltz_proteinfold_output( f"{run.id}/boltz/top_ranked_structures/T1024.pdb", "T1024" ) == ClassifiedOutput("pdb", "T1024.pdb") + assert classify_boltz_proteinfold_output( + f"{run.id}/boltz/top_ranked_structures/T1024.cif", "T1024" + ) == ClassifiedOutput("pdb", "T1024.cif") assert classify_boltz_proteinfold_output( f"{run.id}/boltz/T1024/abcd1234.tsv", "T1024" ) == ClassifiedOutput("stats_csv", "abcd1234.tsv") @@ -474,6 +655,7 @@ def test_boltz_proteinfold_helpers_classify_keys_and_build_prefixes(): ) == ClassifiedOutput("alignment", "single_prediction.a3m") assert build_boltz_proteinfold_output_listing_prefixes(run) == [ + f"{run.id}/", f"{run.id}/reports/", f"{run.id}/boltz/top_ranked_structures/", f"{run.id}/mmseqs/", @@ -530,6 +712,7 @@ def test_alphafold2_proteinfold_helpers_classify_keys_and_build_prefixes(): ) == ClassifiedOutput("stats_csv", "T1024_0_pae.tsv") assert build_alphafold2_proteinfold_output_listing_prefixes(run) == [ + f"{run.id}/", f"{run.id}/reports/", f"{run.id}/alphafold2/split_msa_prediction/top_ranked_structures/", f"{run.id}/alphafold2/split_msa_prediction/T1024/", @@ -589,6 +772,7 @@ def test_colabfold_proteinfold_helpers_classify_keys_and_build_prefixes(): ) == ClassifiedOutput("alignment", "single_prediction.a3m") assert build_colabfold_proteinfold_output_listing_prefixes(run) == [ + f"{run.id}/", f"{run.id}/reports/", f"{run.id}/colabfold/top_ranked_structures/", f"{run.id}/mmseqs/", @@ -881,15 +1065,41 @@ async def test_extract_wisps_max_score_skips_non_numeric_iptm(): assert score == 0.75 +# --------------------------------------------------------------------------- +# get_run_service_usage +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_run_service_usage(): + csv_text = ( + "Name,Queue,Service Units\n" + "prepare,normalbw,0.00\n" + "search,normalbw,0.20\n" + "predict,gpuhopper,3.25\n" + ) + + with patch( + "app.services.results_utils.read_s3_file", + new_callable=AsyncMock, + return_value=csv_text, + ) as mock_read: + service_usage = await get_run_service_usage("run-1/UsageReport.csv") + + assert service_usage == 3.45 + mock_read.assert_awaited_once_with("run-1/UsageReport.csv") + + # --------------------------------------------------------------------------- # build_wisps_output_listing_prefixes # --------------------------------------------------------------------------- -def test_build_wisps_output_listing_prefixes_returns_three_prefixes(): +def test_build_wisps_output_listing_prefixes(): run = WorkflowRun(id=uuid4(), owner_user_id=uuid4(), sample_id="sample1") prefixes = build_wisps_output_listing_prefixes(run) - assert len(prefixes) == 3 + assert len(prefixes) == 4 + assert f"{run.id}/" in prefixes assert f"{run.id}/multiqc/" in prefixes assert f"{run.id}/collect/" in prefixes assert f"{run.id}/ipsae/" in prefixes @@ -914,3 +1124,32 @@ def test_workflow_output_specs_has_interaction_screening(): assert spec.kind == "interaction-screening" assert spec.get_score_file is not None assert spec.extract_max_score is not None + + +def test_classify_shared_outputs_usage_report_csv(): + result = classify_shared_outputs("runs/run-1/UsageReport.csv") + + assert result == ClassifiedOutput(category="usage", label="UsageReport.csv") + + +def test_workflow_results_spec_classify_output_uses_shared_outputs(): + """Test WorkflowResultsSpec classifies shared outputs before workflow-specific outputs.""" + + async def extract_max_score(_key: str) -> float | None: + return None + + classifier = MagicMock(return_value=ClassifiedOutput(category="pdb", label="ignored.pdb")) + spec = WorkflowResultsSpec( + kind="single-prediction", + tool="boltz", + required_categories=set(), + get_prefixes=lambda _run: [], + classifier=classifier, + get_score_file=lambda _keys, _sample_id: None, + extract_max_score=extract_max_score, + ) + + result = spec.classify_output("runs/run-1/UsageReport.csv", sample_id="T1024") + + assert result == ClassifiedOutput(category="usage", label="UsageReport.csv") + classifier.assert_not_called() diff --git a/tests/test_services_s3.py b/tests/test_services_s3.py index 8fa5b41..821ddce 100644 --- a/tests/test_services_s3.py +++ b/tests/test_services_s3.py @@ -17,6 +17,7 @@ get_s3_client, list_s3_files, read_csv_from_s3, + read_s3_bytes, read_s3_file, upload_file_to_s3, ) @@ -265,6 +266,43 @@ async def test_read_s3_file_client_error(mock_env_vars, mock_s3_client): await read_s3_file("results/test/missing.tsv") +@pytest.mark.asyncio +async def test_read_s3_bytes_success(mock_env_vars, mock_s3_client): + """Test reading a bytes object from S3.""" + content = b"rank\tscore\n0\t0.91\n" + mock_response = {"Body": MagicMock()} + mock_response["Body"].read.return_value = content + mock_s3_client.get_object.return_value = mock_response + + result = await read_s3_bytes("results/test/file.tsv") + + assert result == content + mock_s3_client.get_object.assert_called_once_with( + Bucket="test-bucket", + Key="results/test/file.tsv", + ) + + +@pytest.mark.asyncio +async def test_read_s3_bytes_missing_bucket(): + """Test reading a text object without bucket configuration.""" + with patch.dict("os.environ", {}, clear=True): + with pytest.raises(S3ConfigurationError, match="AWS_S3_BUCKET"): + await read_s3_bytes("results/test/file.tsv") + + +@pytest.mark.asyncio +async def test_read_s3_bytes_client_error(mock_env_vars, mock_s3_client): + """Test reading a text object wraps S3 client errors.""" + mock_s3_client.get_object.side_effect = ClientError( + {"Error": {"Code": "NoSuchKey", "Message": "Not found"}}, + "get_object", + ) + + with pytest.raises(S3ServiceError, match="Failed to read file from S3"): + await read_s3_bytes("results/test/missing.tsv") + + @pytest.mark.asyncio async def test_read_csv_from_s3_all_columns(mock_env_vars, mock_s3_client): """Test reading CSV with all columns.""" diff --git a/tests/test_services_seqera.py b/tests/test_services_seqera.py index 6ad158c..fad8757 100644 --- a/tests/test_services_seqera.py +++ b/tests/test_services_seqera.py @@ -19,15 +19,13 @@ launch_bindflow_workflow, prepare_bindflow_workflow, ) -from app.services.seqera import WorkflowExecutorError, WorkflowLaunchResult -from app.services.seqera_errors import SeqeraConfigurationError +from app.services.seqera import WorkflowExecutorError, WorkflowLaunchResult, count_active_workflows +from app.services.seqera_errors import SeqeraAPIError, SeqeraConfigurationError from tests.datagen import AppUserFactory, QueuedJobFactory, WorkflowFactory, WorkflowRunFactory _CONFIG_PATH = "/some/bindflow.config" _USER_DETAILS = WorkflowUserDetails( user_email="test@example.com", - full_name="Test_User", - institute="example.com", ip_address="127.0.0.1", ) @@ -343,3 +341,58 @@ async def test_launch_with_custom_params_text(persistent_models): assert "my_custom_param: 42" in params_text assert "another_param: test" in params_text + + +def _totals_by_status_handler(totals: dict[str, int]): + """Fake /workflow endpoint: returns totalSize for whichever status: was searched.""" + + def _handler(request: httpx.Request) -> httpx.Response: + search = request.url.params.get("search", "") + status = search.removeprefix("status:") + return httpx.Response( + 200, json={"workflows": [], "totalSize": totals.get(status, 0), "hasMore": False} + ) + + return _handler + + +@pytest.mark.asyncio +@respx.mock +async def test_count_active_workflows_sums_running_and_submitted_totals(): + respx.get(url__regex=r".*/workflow(\?.*)?$").mock( + side_effect=_totals_by_status_handler({"RUNNING": 3, "SUBMITTED": 2}) + ) + + assert await count_active_workflows() == 5 + + +@pytest.mark.asyncio +@respx.mock +async def test_count_active_workflows_queries_each_status_with_a_minimal_page(): + route = respx.get(url__regex=r".*/workflow(\?.*)?$").mock( + side_effect=_totals_by_status_handler({}) + ) + + await count_active_workflows() + + assert route.call_count == 2 + searches = {call.request.url.params["search"] for call in route.calls} + assert searches == {"status:RUNNING", "status:SUBMITTED"} + assert all(call.request.url.params["max"] == "1" for call in route.calls) + + +@pytest.mark.asyncio +@respx.mock +async def test_count_active_workflows_no_active_runs(): + respx.get(url__regex=r".*/workflow(\?.*)?$").mock(side_effect=_totals_by_status_handler({})) + + assert await count_active_workflows() == 0 + + +@pytest.mark.asyncio +@respx.mock +async def test_count_active_workflows_raises_on_api_error(): + respx.get(url__regex=r".*/workflow(\?.*)?$").mock(return_value=httpx.Response(500, text="boom")) + + with pytest.raises(SeqeraAPIError): + await count_active_workflows() diff --git a/tests/test_services_seqera_client.py b/tests/test_services_seqera_client.py index 243374d..9ffd1fc 100644 --- a/tests/test_services_seqera_client.py +++ b/tests/test_services_seqera_client.py @@ -91,6 +91,47 @@ async def test_describe_and_list_success(monkeypatch): assert await get_workflow_logs_raw("wf-1") == {"ok": True} +@pytest.mark.asyncio +async def test_list_workflows_raw_passes_max_and_search_params(monkeypatch): + monkeypatch.setenv("SEQERA_API_URL", "https://api.seqera.test") + monkeypatch.setenv("SEQERA_ACCESS_TOKEN", "token") + monkeypatch.setenv("WORK_SPACE", "ws-1") + + ok = AsyncMock(spec=httpx.Response) + ok.is_error = False + ok.json.return_value = {"workflows": []} + + with patch("httpx.AsyncClient.get", return_value=ok) as mock_get: + result = await list_workflows_raw(search_query="status:RUNNING", max_results=100) + + assert result == {"workflows": []} + mock_get.assert_awaited_once_with( + "https://api.seqera.test/workflow", + params={"workspaceId": "ws-1", "search": "status:RUNNING", "max": "100"}, + headers={"Authorization": "Bearer token", "Accept": "application/json"}, + ) + + +@pytest.mark.asyncio +async def test_list_workflows_raw_omits_max_when_not_given(monkeypatch): + monkeypatch.setenv("SEQERA_API_URL", "https://api.seqera.test") + monkeypatch.setenv("SEQERA_ACCESS_TOKEN", "token") + monkeypatch.delenv("WORK_SPACE", raising=False) + + ok = AsyncMock(spec=httpx.Response) + ok.is_error = False + ok.json.return_value = {"workflows": []} + + with patch("httpx.AsyncClient.get", return_value=ok) as mock_get: + await list_workflows_raw() + + mock_get.assert_awaited_once_with( + "https://api.seqera.test/workflow", + params={}, + headers={"Authorization": "Bearer token", "Accept": "application/json"}, + ) + + @pytest.mark.asyncio async def test_cancel_and_delete_paths(monkeypatch): monkeypatch.setenv("SEQERA_API_URL", "https://api.seqera.test") diff --git a/tests/test_wisps_coverage.py b/tests/test_wisps_coverage.py index d269ced..dedf83a 100644 --- a/tests/test_wisps_coverage.py +++ b/tests/test_wisps_coverage.py @@ -41,8 +41,6 @@ _USER_DETAILS = WorkflowUserDetails( user_email="user@ex.com", - full_name="Test User", - institute="USYD", ip_address="1.2.3.4", ) @@ -258,33 +256,37 @@ def test_get_wisps_config_text_appends_process_block(): with patch("builtins.open", mock_open(read_data="base_config")): result = get_wisps_config_text( config_file_path="/fake/path.config", - job_id="my-job", user_details=_USER_DETAILS, - timestamp="20240101_120000", ) assert "process {" in result assert "clusterOptions" in result -def test_get_wisps_config_text_contains_job_fields(): +def test_get_wisps_config_text_contains_encoded_email_and_encoded_ip(): with patch("builtins.open", mock_open(read_data="base_config")): result = get_wisps_config_text( config_file_path="/fake/path.config", - job_id="my-job", user_details=_USER_DETAILS, - timestamp="20240101_120000", ) - assert "my-job" in result - assert "user@ex.com" in result + assert "dXNlckBleC5jb20=" in result + assert "MS4yLjMuNA==" in result + + +def test_get_wisps_config_text_without_ip_uses_encoded_email_only(): + with patch("builtins.open", mock_open(read_data="base_config")): + result = get_wisps_config_text( + config_file_path="/fake/path.config", + user_details=_USER_DETAILS.model_copy(update={"ip_address": ""}), + ) + assert "-A dXNlckBleC5jb20=" in result + assert ":" not in result.split("clusterOptions = ")[1] def test_get_wisps_config_text_contains_base_config(): with patch("builtins.open", mock_open(read_data="base_config")): result = get_wisps_config_text( config_file_path="/fake/path.config", - job_id="my-job", user_details=_USER_DETAILS, - timestamp="20240101_120000", ) assert "base_config" in result diff --git a/uv.lock b/uv.lock index 14d44fb..1d6be17 100644 --- a/uv.lock +++ b/uv.lock @@ -397,6 +397,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, ] +[[package]] +name = "httpcore2" +version = "2.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "truststore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/39/a8/20ed1ed79cbc2ecdf5301c0968ab7c85547212e2a7bd126ddd2d986e206e/httpcore2-2.9.1.tar.gz", hash = "sha256:4d8acbf8b306f48c9d6046591fd5ba4037d1b1b1000d140fc2c3eab1e9a0c0e2", size = 67089, upload-time = "2026-07-24T09:21:03.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/fb/46c52b781975c335a2bcf1072c7bbc007cbdc8d674217f5ee1daba2c848b/httpcore2-2.9.1-py3-none-any.whl", hash = "sha256:6182472379e855fe4221246a2bb7ecede403bc61c6798062ae1787d051ccde26", size = 82809, upload-time = "2026-07-24T09:21:01.178Z" }, +] + [[package]] name = "httptools" version = "0.8.0" @@ -434,6 +447,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, ] +[[package]] +name = "httpx2" +version = "2.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "httpcore2" }, + { name = "idna" }, + { name = "truststore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/21/14/38128fbafd7e0ed41d874df6c9a653d47c2d111cfe59e2b4ac95161b4abd/httpx2-2.9.1.tar.gz", hash = "sha256:1932a768737e3666291582833da748cc4e563c337cf96706fccc04fa6e58764a", size = 95458, upload-time = "2026-07-24T09:21:04.972Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/b8/cfd91c4ab9134d386d48f0b6ac662ff3d4be6efdee59ee1c67ebc3c0487c/httpx2-2.9.1-py3-none-any.whl", hash = "sha256:1820fe14a9ab1107bfeff39259987429450b070ec0ff38cc87eb0d8c97fdc71a", size = 91191, upload-time = "2026-07-24T09:21:02.6Z" }, +] + [[package]] name = "idna" version = "3.18" @@ -1083,7 +1111,6 @@ dependencies = [ { name = "sqlalchemy" }, { name = "starlette-admin" }, { name = "typer" }, - { name = "unidecode" }, { name = "uvicorn", extra = ["standard"] }, ] @@ -1096,6 +1123,7 @@ dev = [ { name = "faker" }, { name = "groovy-parser" }, { name = "httpx" }, + { name = "httpx2" }, { name = "mypy" }, { name = "polyfactory" }, { name = "pytest" }, @@ -1126,7 +1154,6 @@ requires-dist = [ { name = "sqlalchemy", specifier = "~=2.0" }, { name = "starlette-admin", specifier = "~=0.16" }, { name = "typer", specifier = ">=0.26.8" }, - { name = "unidecode", specifier = "~=1.3" }, { name = "uvicorn", extras = ["standard"], specifier = "~=0.49" }, ] @@ -1139,6 +1166,7 @@ dev = [ { name = "faker", specifier = ">=40.1.0" }, { name = "groovy-parser", git = "https://github.com/inab/python-groovy-parser.git?rev=5e7f3a4e7dcc5fb2e186fe7ca30df260ce5efeb8" }, { name = "httpx", specifier = "~=0.28" }, + { name = "httpx2", specifier = "~=2.9" }, { name = "mypy", specifier = "~=1.19" }, { name = "polyfactory", specifier = "~=2.16" }, { name = "pytest", specifier = "~=8.3" }, @@ -1236,6 +1264,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/7c/add7d4be6690fb61f4d30b672e544b1ac541c824f05a7de50ff5d8985db5/starlette_admin-0.16.1-py3-none-any.whl", hash = "sha256:a5e6cb0beb2e9bdc65b07dbb4bd01726aaad34764595e3bc3ff8578e698ff98a", size = 2176514, upload-time = "2026-06-06T20:10:42.308Z" }, ] +[[package]] +name = "truststore" +version = "0.10.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/a3/1585216310e344e8102c22482f6060c7a6ea0322b63e026372e6dcefcfd6/truststore-0.10.4.tar.gz", hash = "sha256:9d91bd436463ad5e4ee4aba766628dd6cd7010cf3e2461756b3303710eebc301", size = 26169, upload-time = "2025-08-12T18:49:02.73Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/97/56608b2249fe206a67cd573bc93cd9896e1efb9e98bce9c163bcdc704b88/truststore-0.10.4-py3-none-any.whl", hash = "sha256:adaeaecf1cbb5f4de3b1959b42d41f6fab57b2b1666adb59e89cb0b53361d981", size = 18660, upload-time = "2025-08-12T18:49:01.46Z" }, +] + [[package]] name = "typer" version = "0.26.8" @@ -1320,15 +1357,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9e/a4/017a7a6cbe387d961a688ec31364ae60a5c4e22c96ae9921b79a947c855d/tzlocal-5.4.4-py3-none-any.whl", hash = "sha256:aae09f0126a8a86fa736be266eb4a471380d26a0de3bc14844e7821fee3e2a15", size = 18115, upload-time = "2026-06-29T08:03:38.666Z" }, ] -[[package]] -name = "unidecode" -version = "1.4.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/94/7d/a8a765761bbc0c836e397a2e48d498305a865b70a8600fd7a942e85dcf63/Unidecode-1.4.0.tar.gz", hash = "sha256:ce35985008338b676573023acc382d62c264f307c8f7963733405add37ea2b23", size = 200149, upload-time = "2025-04-24T08:45:03.798Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/b7/559f59d57d18b44c6d1250d2eeaa676e028b9c527431f5d0736478a73ba1/Unidecode-1.4.0-py3-none-any.whl", hash = "sha256:c3c7606c27503ad8d501270406e345ddb480a7b5f38827eafe4fa82a137f0021", size = 235837, upload-time = "2025-04-24T08:45:01.609Z" }, -] - [[package]] name = "urllib3" version = "2.7.0"