feat: track service unit cost of runs in database (SBP-445)

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 ###

app/db/admin.py

@@ -168,8 +168,9 @@ class WorkflowRunAdmin(ModelView):
         HasOne("owner", identity="app-user"),
         "seqera_run_id",
         "run_name",
-        "binder_name",
+        "service_usage",
         JSONField("submitted_form_data"),
+        "binder_name",
         "work_dir",
         "submission_timestamp",
     ]

app/db/models/core.py

@@ -7,6 +7,7 @@
     JSON,
     BigInteger,
     DateTime,
+    Float,
     ForeignKey,
     Index,
     Numeric,
@@ -92,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")

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:

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

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"
@@ -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/",
@@ -697,7 +750,7 @@ 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}/results/"]
+    return [f"{run_uuid}/", f"{run_uuid}/results/"]
 
 
 def _make_wisps_spec(tool: WorkflowTool) -> WorkflowResultsSpec:
@@ -708,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,
     )
 
 
@@ -720,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,
     )
 
 
@@ -733,7 +786,7 @@ 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(
@@ -743,7 +796,7 @@ def _make_bulk_prediction_spec(tool: WorkflowTool) -> WorkflowResultsSpec:
             get_prefixes=build_rfdiffusion_output_listing_prefixes,
             get_score_file=get_rfdiffusion_score_file,
             extract_max_score=extract_rfdiffusion_max_score,
-            classify=classify_rfdiffusion_output_key,
+            classifier=classify_rfdiffusion_output_key,
         ),
     },
     "single-prediction": {
@@ -754,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",
@@ -763,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",
@@ -772,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": {
@@ -802,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
@@ -881,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
 
@@ -966,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(
@@ -990,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()
@@ -1083,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)

pyproject.toml

@@ -41,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",

tests/conftest.py

@@ -140,6 +140,7 @@ def persistent_models(test_db):
         datagen.WorkflowRunFactory,
         datagen.RunInputFactory,
         datagen.RunOutputFactory,
+        datagen.S3ObjectFactory,
         datagen.QueuedJobFactory,
     ]
 

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