From b20aeef8f75fc5f7401399bd59a6afec66665ae4 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 15:31:27 +1000 Subject: [PATCH 01/16] feat: add service_usage column to WorkflowRun --- ...60728_145217_service_usage_4f26ab5038ec.py | 23 +++++++++++++++++++ app/db/models/core.py | 2 ++ 2 files changed, 25 insertions(+) create mode 100644 alembic/versions/20260728_145217_service_usage_4f26ab5038ec.py 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/models/core.py b/app/db/models/core.py index 8440057..6ca953b 100644 --- a/app/db/models/core.py +++ b/app/db/models/core.py @@ -14,6 +14,7 @@ Text, UniqueConstraint, text, + Float, ) from sqlalchemy.dialects.postgresql import INET, UUID from sqlalchemy.orm import Mapped, Session, mapped_column, relationship @@ -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") From 10591698a98cb240deb2d67746152d46b11d2578 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 15:43:02 +1000 Subject: [PATCH 02/16] feat: function to sync service usage for a run --- app/services/job_utils.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/app/services/job_utils.py b/app/services/job_utils.py index 3015d95..f46cefa 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: + 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 From d97e0775ece72ecd985e0dba6b938c0991bcc455 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 16:16:55 +1000 Subject: [PATCH 03/16] fix: check current service usage is not None --- app/services/job_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/services/job_utils.py b/app/services/job_utils.py index f46cefa..c01e775 100644 --- a/app/services/job_utils.py +++ b/app/services/job_utils.py @@ -211,7 +211,7 @@ async def ensure_completed_run_score(db: Session, run: WorkflowRun, ui_status: s async def sync_service_usage(db: Session, run: WorkflowRun, ui_status: str) -> float | None: if ui_status != "Completed": return None - if run.service_usage: + if run.service_usage is not None: return run.service_usage try: From 0c00f39998ad1defe15569bf5b62787796df8850 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 16:18:28 +1000 Subject: [PATCH 04/16] refactor: rework workflow results so we can get service usage --- app/services/results_utils.py | 88 ++++++++++++++++++++++++++++++----- 1 file changed, 76 insertions(+), 12 deletions(-) diff --git a/app/services/results_utils.py b/app/services/results_utils.py index 93c24e4..3b08061 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,21 @@ 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 +558,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 +583,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 +605,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 +677,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 +692,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/", @@ -662,7 +714,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 +726,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,7 +739,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, ), }, @@ -699,7 +751,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 +760,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 +769,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 +799,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 +878,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 +963,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 +987,8 @@ 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 +1081,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) From 6a980b081f06fdee96f25ac35473f667b78015b7 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 16:19:23 +1000 Subject: [PATCH 05/16] feat: sync service usage when looking up jobs (current syncs trigger here) --- app/routes/workflow/jobs.py | 8 ++++++++ 1 file changed, 8 insertions(+) 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: From 0862e250f41fecfbd8ec37cf7e6f44efebe3fe96 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 16:25:41 +1000 Subject: [PATCH 06/16] test: update tests with reworked results/outputs --- tests/test_routes_workflow_jobs_extra.py | 4 ++-- tests/test_services_results_utils.py | 14 ++++++++++---- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/test_routes_workflow_jobs_extra.py b/tests/test_routes_workflow_jobs_extra.py index 959ae81..4ed1537 100644 --- a/tests/test_routes_workflow_jobs_extra.py +++ b/tests/test_routes_workflow_jobs_extra.py @@ -96,7 +96,7 @@ async def test_get_job_details_success(test_db): ) workflow = Workflow( id=uuid4(), - name="BindCraft", + name="de-novo-design", description="Binding workflow", repo_url="https://github.com/test/bindcraft", default_revision="main", @@ -136,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_services_results_utils.py b/tests/test_services_results_utils.py index 714054d..8f5c4e9 100644 --- a/tests/test_services_results_utils.py +++ b/tests/test_services_results_utils.py @@ -210,6 +210,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 +218,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/", ] @@ -374,7 +376,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, ) @@ -420,7 +422,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, ) @@ -482,6 +484,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/", @@ -538,6 +541,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/", @@ -597,6 +601,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/", @@ -894,10 +899,11 @@ async def test_extract_wisps_max_score_skips_non_numeric_iptm(): # --------------------------------------------------------------------------- -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 From fda667c6a9f0e3de15f3b2375c12203cf388788c Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 16:36:24 +1000 Subject: [PATCH 07/16] chore: add httpx2, recommended by FastAPI/starlette --- pyproject.toml | 1 + uv.lock | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 91b0591..67a813c 100644 --- a/pyproject.toml +++ b/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", diff --git a/uv.lock b/uv.lock index ec6b5d1..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" @@ -1095,6 +1123,7 @@ dev = [ { name = "faker" }, { name = "groovy-parser" }, { name = "httpx" }, + { name = "httpx2" }, { name = "mypy" }, { name = "polyfactory" }, { name = "pytest" }, @@ -1137,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" }, @@ -1234,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" From 47bf33e491a2d21d6443a82799dd7e0d765bd4ef Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Tue, 28 Jul 2026 16:36:51 +1000 Subject: [PATCH 08/16] test: unit test of service usage calculation --- tests/test_services_results_utils.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/test_services_results_utils.py b/tests/test_services_results_utils.py index 8f5c4e9..e524a2b 100644 --- a/tests/test_services_results_utils.py +++ b/tests/test_services_results_utils.py @@ -33,6 +33,7 @@ get_all_downloads_zipped, get_bindcraft_score_file, get_proteinfold_score_file, + get_run_service_usage, get_sample_id_for_result, get_tool_name, get_wisps_score_file, @@ -894,6 +895,31 @@ 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 # --------------------------------------------------------------------------- From bbbf09c56933eb87abb105f6490e32acd05b1e65 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 09:50:13 +1000 Subject: [PATCH 09/16] test: add test factory for S3Object --- tests/conftest.py | 1 + tests/datagen.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index e1f0a0b..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, ] 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 From fd2df572817fdab706d853fe07756558d0e8b0df Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 09:53:37 +1000 Subject: [PATCH 10/16] test: add test of extracting service units for run --- tests/test_services_results_utils.py | 46 +++++++++++++++++++++++++++- 1 file changed, 45 insertions(+), 1 deletion(-) diff --git a/tests/test_services_results_utils.py b/tests/test_services_results_utils.py index e524a2b..c96b32a 100644 --- a/tests/test_services_results_utils.py +++ b/tests/test_services_results_utils.py @@ -43,7 +43,7 @@ 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(): @@ -432,6 +432,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") From 4315891ec1548819105c6b5ef341213b0e44d3f4 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 11:11:07 +1000 Subject: [PATCH 11/16] fix: update arg name in rfdiffusion spec --- app/services/results_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/services/results_utils.py b/app/services/results_utils.py index 9f91611..7e93974 100644 --- a/app/services/results_utils.py +++ b/app/services/results_utils.py @@ -795,7 +795,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": { From b308db1c5859f0f8bd6f8a1f916a23fb5147460e Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 11:16:53 +1000 Subject: [PATCH 12/16] fix: linting --- app/db/models/core.py | 2 +- app/services/results_utils.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/app/db/models/core.py b/app/db/models/core.py index 6ca953b..ac5736f 100644 --- a/app/db/models/core.py +++ b/app/db/models/core.py @@ -7,6 +7,7 @@ JSON, BigInteger, DateTime, + Float, ForeignKey, Index, Numeric, @@ -14,7 +15,6 @@ Text, UniqueConstraint, text, - Float, ) from sqlalchemy.dialects.postgresql import INET, UUID from sqlalchemy.orm import Mapped, Session, mapped_column, relationship diff --git a/app/services/results_utils.py b/app/services/results_utils.py index 7e93974..3ba486c 100644 --- a/app/services/results_utils.py +++ b/app/services/results_utils.py @@ -135,6 +135,7 @@ def classify_shared_outputs(key: str) -> ClassifiedOutput | None: return None + def _sanitize_content_disposition_filename(filename: str) -> str: sanitized = _HEADER_UNSAFE_FILENAME_CHARS.sub("_", filename).strip() return sanitized or "download" @@ -1042,8 +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 not in ("snapshot", "usage")] + downloads = [ + (key, output) + for key, output in outputs.items() + if output.category not in ("snapshot", "usage") + ] used_filenames: set[str] = set() zip_file = BytesIO() From ad497cb2bdb00e8dc8790945d2b954fdfe73e108 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 13:02:14 +1000 Subject: [PATCH 13/16] test: unit tests of read_s3_bytes --- tests/test_services_s3.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_services_s3.py b/tests/test_services_s3.py index 8fa5b41..148ae66 100644 --- a/tests/test_services_s3.py +++ b/tests/test_services_s3.py @@ -19,6 +19,7 @@ read_csv_from_s3, read_s3_file, upload_file_to_s3, + read_s3_bytes, ) @@ -265,6 +266,44 @@ 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.""" From ae13ce12756898921f32850d654784e8fba1a9eb Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 13:25:35 +1000 Subject: [PATCH 14/16] test: more tests of result_utils to improve coverage --- app/services/results_utils.py | 2 +- tests/test_services_results_utils.py | 157 ++++++++++++++++++++++++++- 2 files changed, 157 insertions(+), 2 deletions(-) diff --git a/app/services/results_utils.py b/app/services/results_utils.py index 3ba486c..0f98c7b 100644 --- a/app/services/results_utils.py +++ b/app/services/results_utils.py @@ -750,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: diff --git a/tests/test_services_results_utils.py b/tests/test_services_results_utils.py index c96b32a..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,23 +20,29 @@ 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, @@ -230,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 @@ -355,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( @@ -998,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() From fab52011cac249a57ef93cdc8e2442320c22f1b9 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 13:37:34 +1000 Subject: [PATCH 15/16] chore: linting --- tests/test_services_s3.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_services_s3.py b/tests/test_services_s3.py index 148ae66..821ddce 100644 --- a/tests/test_services_s3.py +++ b/tests/test_services_s3.py @@ -17,9 +17,9 @@ get_s3_client, list_s3_files, read_csv_from_s3, + read_s3_bytes, read_s3_file, upload_file_to_s3, - read_s3_bytes, ) @@ -303,7 +303,6 @@ async def test_read_s3_bytes_client_error(mock_env_vars, mock_s3_client): 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.""" From b09bef885840bbe593ffa4ece9812d466a644620 Mon Sep 17 00:00:00 2001 From: Marius Mather Date: Wed, 29 Jul 2026 13:48:22 +1000 Subject: [PATCH 16/16] chore: add service_usage to workflowrun admin --- app/db/admin.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/app/db/admin.py b/app/db/admin.py index 04fad8a..6326298 100644 --- a/app/db/admin.py +++ b/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", ]