From 1b71f8466b466da592a26f6ae690ddff5aa5a2f9 Mon Sep 17 00:00:00 2001 From: Mihir Samdarshi Date: Sat, 4 Apr 2026 14:20:12 +0530 Subject: [PATCH] feat: add GCP cost estimation support - Add CromwellBackendCostService for GCP cost configuration - Extend API with cost endpoint and estimation retrieval - Add client-side cost command with JSON/TSV/human output formats - Support configurable catalog expiry for pricing data --- caper/caper_args.py | 50 ++++++++++ caper/caper_backend_conf.py | 17 ++++ caper/caper_client.py | 18 ++++ caper/caper_init.py | 4 + caper/caper_runner.py | 9 ++ caper/cli.py | 56 +++++++++++ caper/cromwell_backend.py | 24 +++++ caper/cromwell_rest_api.py | 27 +++++ tests/test_cost.py | 176 +++++++++++++++++++++++++++++++++ tests/test_cromwell_backend.py | 26 ++++- 10 files changed, 405 insertions(+), 2 deletions(-) create mode 100644 tests/test_cost.py diff --git a/caper/caper_args.py b/caper/caper_args.py index 015cd7b..cb05679 100644 --- a/caper/caper_args.py +++ b/caper/caper_args.py @@ -354,6 +354,25 @@ def _add_gcp_runner_args(parser: ArgumentParser) -> None: default='mirror.gcr.io', help='Address of the Docker Hub mirror (default: mirror.gcr.io).', ) + group_gc.add_argument( + '--gcp-cost-estimation', + action=argparse.BooleanOptionalAction, + default=False, + help=( + 'Enable GCP cost estimation in Cromwell server mode. ' + 'Requires Cloud Billing API enabled on your GCP project. ' + 'Use --no-gcp-cost-estimation to disable.' + ), + ) + group_gc.add_argument( + '--gcp-cost-catalog-expiry', + type=int, + default=86400, + help=( + 'How long (seconds) Cromwell caches the GCP pricing catalog ' + 'before re-fetching from Cloud Billing API (default: 86400).' + ), + ) def _add_aws_runner_args(parser: ArgumentParser) -> None: @@ -891,6 +910,23 @@ def _create_parent_cleanup() -> ArgumentParser: return parser +def _add_cost_args(parser: ArgumentParser) -> None: + group = parser.add_argument_group(title='cost command arguments') + group.add_argument( + '--format', + dest='cost_format', + choices=['human', 'tsv', 'json'], + default='human', + help='Output format for cost results (default: human).', + ) + + +def _create_parent_cost() -> ArgumentParser: + parser = argparse.ArgumentParser(add_help=False) + _add_cost_args(parser) + return parser + + def _create_parent_hpc_abort() -> ArgumentParser: parser = argparse.ArgumentParser(add_help=False) _add_hpc_abort_args(parser) @@ -964,6 +1000,7 @@ def get_parser_and_defaults( parent_gcp_res_analysis = _create_parent_gcp_res_analysis() parent_cleanup = _create_parent_cleanup() parent_hpc_abort = _create_parent_hpc_abort() + parent_cost = _create_parent_cost() # all subcommands p_init = subparser.add_parser( @@ -1050,6 +1087,18 @@ def get_parser_and_defaults( ], ) + p_cost = subparser.add_parser( + 'cost', + help='Get estimated GCP cost for workflows (server mode only).', + parents=[ + parent_all, + parent_server_client, + parent_client, + parent_search_wf, + parent_cost, + ], + ) + p_hpc = subparser.add_parser('hpc', help='Subcommand for HPCs', parents=[parent_all]) subparser_hpc = p_hpc.add_subparsers(dest='hpc_action') subparser_hpc.add_parser( @@ -1135,6 +1184,7 @@ def get_parser_and_defaults( p_metadata, p_troubleshoot, p_debug, + p_cost, p_gcp_monitor, p_gcp_res_analysis, p_cleanup, diff --git a/caper/caper_backend_conf.py b/caper/caper_backend_conf.py index 0254bc4..f11fdd3 100644 --- a/caper/caper_backend_conf.py +++ b/caper/caper_backend_conf.py @@ -1,4 +1,5 @@ """Module for creating the Cromwell backend configuration file.""" + from __future__ import annotations import logging @@ -14,6 +15,7 @@ CromwellBackendAws, CromwellBackendBase, CromwellBackendCommon, + CromwellBackendCostService, CromwellBackendDatabase, CromwellBackendGcp, CromwellBackendLocal, @@ -71,6 +73,8 @@ def __init__( gcp_dockerhub_mirror: bool = True, gcp_dockerhub_mirror_address: str = 'mirror.gcr.io', gcp_region: str = CromwellBackendGcp.DEFAULT_REGION, + gcp_cost_estimation: bool = False, + gcp_cost_catalog_expiry: int = CromwellBackendCostService.DEFAULT_CATALOG_EXPIRY_SECONDS, aws_batch_arn: str | None = None, aws_region: str | None = None, aws_out_dir: str | None = None, @@ -171,6 +175,10 @@ def __init__( Address of the Docker Hub mirror. gcp_region: Region for Google Cloud Batch API. + gcp_cost_estimation: + Enable GCP cost estimation service in Cromwell. + gcp_cost_catalog_expiry: + How long (seconds) Cromwell caches the GCP pricing catalog. aws_batch_arn: ARN for AWS Batch. aws_region: @@ -335,6 +343,15 @@ def __init__( ), ) + if gcp_cost_estimation: + merge_dict( + self._template, + CromwellBackendCostService( + enabled=True, + catalog_expiry_seconds=gcp_cost_catalog_expiry, + ), + ) + if aws_batch_arn and aws_region and aws_out_dir: merge_dict( self._template, diff --git a/caper/caper_client.py b/caper/caper_client.py index 9a7be9a..55c8cbd 100644 --- a/caper/caper_client.py +++ b/caper/caper_client.py @@ -153,6 +153,24 @@ def metadata( workflow_ids, labels, embed_subworkflow=embed_subworkflow ) + def cost( + self, + wf_ids_or_labels: Sequence[str], + ) -> list[dict[str, Any]] | None: + """Retrieve estimated costs for workflows from a Cromwell server. + + Args: + wf_ids_or_labels: + List of workflow IDs or string labels (Caper's string label). + Wildcards (*, ?) are allowed. + + Returns: + List of cost response dicts for matched workflows. + """ + workflow_ids, labels = self._split_workflow_ids_and_labels(wf_ids_or_labels) + + return self._cromwell_rest_api.get_cost(workflow_ids, labels) + def _split_workflow_ids_and_labels( self, workflow_ids_or_labels: Sequence[str] | None ) -> tuple[Sequence[str], Sequence[tuple[str, str]]]: diff --git a/caper/caper_init.py b/caper/caper_init.py index f30d82c..397d7f7 100644 --- a/caper/caper_init.py +++ b/caper/caper_init.py @@ -110,6 +110,10 @@ # e.g. us-central1 gcp-region= +# Enable GCP cost estimation (requires Cloud Billing API enabled on your project) +# See: https://cromwell.readthedocs.io/en/stable/cromwell_features/CostEstimation/ +gcp-cost-estimation= + # Number of retrials. This parameter also applies to non-OOM failures. max-retries=1 """ diff --git a/caper/caper_runner.py b/caper/caper_runner.py index 75f1e23..0f0ecba 100644 --- a/caper/caper_runner.py +++ b/caper/caper_runner.py @@ -17,6 +17,7 @@ CromwellBackendAws, CromwellBackendBase, CromwellBackendCommon, + CromwellBackendCostService, CromwellBackendDatabase, CromwellBackendGcp, CromwellBackendLocal, @@ -85,6 +86,8 @@ def __init__( gcp_dockerhub_mirror: bool = True, gcp_dockerhub_mirror_address: str = 'mirror.gcr.io', gcp_region: str = CromwellBackendGcp.DEFAULT_REGION, + gcp_cost_estimation: bool = False, + gcp_cost_catalog_expiry: int = CromwellBackendCostService.DEFAULT_CATALOG_EXPIRY_SECONDS, aws_batch_arn: str | None = None, aws_region: str | None = None, aws_out_dir: str | None = None, @@ -186,6 +189,10 @@ def __init__( Address of the Docker Hub mirror (default: mirror.gcr.io). gcp_region: Region for Google Cloud Batch API. + gcp_cost_estimation: + Enable GCP cost estimation service in Cromwell. + gcp_cost_catalog_expiry: + How long (seconds) Cromwell caches the GCP pricing catalog. aws_batch_arn: ARN for AWS Batch. aws_region: @@ -271,6 +278,8 @@ def __init__( gcp_dockerhub_mirror=gcp_dockerhub_mirror, gcp_dockerhub_mirror_address=gcp_dockerhub_mirror_address, gcp_region=gcp_region, + gcp_cost_estimation=gcp_cost_estimation, + gcp_cost_catalog_expiry=gcp_cost_catalog_expiry, aws_batch_arn=aws_batch_arn, aws_region=aws_region, aws_out_dir=aws_out_dir, diff --git a/caper/cli.py b/caper/cli.py index 5b4d4e7..a7de81c 100644 --- a/caper/cli.py +++ b/caper/cli.py @@ -228,6 +228,8 @@ def runner(args: Namespace, nonblocking_server: bool = False) -> NBSubprocThread postgresql_db_name=args.postgresql_db_name, gcp_prj=args.gcp_prj, gcp_region=args.gcp_region, + gcp_cost_estimation=args.action == 'server' and args.gcp_cost_estimation, + gcp_cost_catalog_expiry=args.gcp_cost_catalog_expiry, gcp_zones=args.gcp_zones, gcp_call_caching_dup_strat=args.gcp_call_caching_dup_strat, gcp_out_dir=args.gcp_out_dir, @@ -321,6 +323,8 @@ def client(args: Namespace) -> None: subcmd_gcp_res_analysis(c, args) elif args.action == 'cleanup': subcmd_cleanup(c, args) + elif args.action == 'cost': + subcmd_cost(c, args) else: msg = f'Unsupported client action {args.action}' raise ValueError(msg) @@ -522,6 +526,58 @@ def subcmd_metadata(caper_client: CaperClient, args: Namespace) -> None: print(json.dumps(m[0], indent=4)) # noqa: T201 +def _format_cost_human(results: list[dict[str, Any]]) -> None: + """Print cost results in human-readable format.""" + has_warnings = False + total_cost = 0.0 + for r in results: + cost = r['cost'] + total_cost += cost + status = r['status'] + wf_id = r['id'] + print(f' {wf_id} {status:<12} ${cost:.2f}') # noqa: T201 + if r.get('errors'): + has_warnings = True + + if len(results) > 1 or has_warnings: + print(' ' + '\u2500' * 60) # noqa: T201 + label = f'Total ({len(results)} workflows' + if has_warnings: + label += ', incomplete' + label += ')' + print(f' {label:<50} ${total_cost:.2f}') # noqa: T201 + + if has_warnings: + print() # noqa: T201 + print('Warnings:') # noqa: T201 + for r in results: + for err in r.get('errors', []): + short_id = r['id'][:8] + print(f' {short_id}...: {err}') # noqa: T201 + + +def subcmd_cost(caper_client: CaperClient, args: Namespace) -> None: + """Execute the cost subcommand.""" + results = caper_client.cost(wf_ids_or_labels=args.wf_id_or_label) + if not results: + msg = 'Found no workflow matching with search query.' + raise ValueError(msg) + + fmt = args.cost_format + + if fmt == 'json': + print(json.dumps(results, indent=4)) # noqa: T201 + elif fmt == 'tsv': + print('workflow_id\tstatus\tcost\tcurrency\terrors') # noqa: T201 + for r in results: + errors_str = '; '.join(r.get('errors', [])) + print( # noqa: T201 + f'{r["id"]}\t{r["status"]}\t{r["cost"]}\t{r["currency"]}\t{errors_str}' + ) + else: + _format_cost_human(results) + + def get_single_cromwell_metadata_obj( caper_client: CaperClient, args: Namespace, subcmd: str ) -> CromwellMetadata: diff --git a/caper/cromwell_backend.py b/caper/cromwell_backend.py index cf9f8d2..059e615 100644 --- a/caper/cromwell_backend.py +++ b/caper/cromwell_backend.py @@ -259,6 +259,30 @@ def __init__( raise ValueError(msg) +class CromwellBackendCostService(UserDict): + """Cromwell GCP cost estimation service configuration.""" + + DEFAULT_CATALOG_EXPIRY_SECONDS = 86400 + + def __init__( + self, + *, + enabled: bool = False, + catalog_expiry_seconds: int = DEFAULT_CATALOG_EXPIRY_SECONDS, + ) -> None: + super().__init__() + if enabled: + self['services'] = { + 'GcpCostCatalogService': { + 'class': 'cromwell.services.cost.GcpCostCatalogService', + 'config': { + 'enabled': True, + 'catalogExpirySeconds': catalog_expiry_seconds, + }, + } + } + + class CromwellBackendBase(UserDict): """Base skeleton backend for all backends.""" diff --git a/caper/cromwell_rest_api.py b/caper/cromwell_rest_api.py index 9ec8ebd..7e13e94 100644 --- a/caper/cromwell_rest_api.py +++ b/caper/cromwell_rest_api.py @@ -99,6 +99,7 @@ class CromwellRestAPI: ENDPOINT_SUBMIT = '/api/workflows/v1' ENDPOINT_ABORT = '/api/workflows/v1/{wf_id}/abort' ENDPOINT_RELEASE_HOLD = '/api/workflows/v1/{wf_id}/releaseHold' + ENDPOINT_COST = '/api/workflows/v1/{wf_id}/cost' DEFAULT_HOSTNAME = 'localhost' DEFAULT_PORT = 8000 @@ -303,6 +304,32 @@ def get_metadata( result.append(cm.metadata) return result + def get_cost( + self, + workflow_ids: list[str] | None = None, + labels: dict[str, str] | None = None, + ) -> list[dict] | None: + """Retrieve estimated cost for workflows matching workflow IDs or labels. + + Args: + workflow_ids: + List of workflow IDs to find workflows matched. + labels: + List of Caper's string labels to find workflows matched. + """ + valid_workflow_ids = self.find_valid_workflow_ids(workflow_ids=workflow_ids, labels=labels) + if valid_workflow_ids is None: + return None + + result = [] + for workflow_id in valid_workflow_ids: + r = self.__request_get( + CromwellRestAPI.ENDPOINT_COST.format(wf_id=workflow_id), + ) + if r: + result.append(r) + return result + def get_labels(self, workflow_id: str) -> dict | None: """ Get labels JSON for a specified workflow. diff --git a/tests/test_cost.py b/tests/test_cost.py new file mode 100644 index 0000000..9b8017c --- /dev/null +++ b/tests/test_cost.py @@ -0,0 +1,176 @@ +"""Tests for GCP cost estimation feature.""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from caper.caper_client import CaperClient +from caper.cli import subcmd_cost +from caper.cromwell_rest_api import CromwellRestAPI + + +SAMPLE_COST_RESPONSE = { + 'id': 'f9c26f2e-f550-4748-a650-5d0d4cab9f3a', + 'status': 'Succeeded', + 'cost': 1.23, + 'currency': 'USD', + 'errors': [], +} + +SAMPLE_COST_RESPONSE_WITH_ERRORS = { + 'id': 'a1b2c3d4-e5f6-7890-abcd-ef1234567890', + 'status': 'Running', + 'cost': 0.45, + 'currency': 'USD', + 'errors': ["Couldn't find valid vmCostPerHour for callA.-1.1"], +} + + +class TestGetCost: + """Tests for CromwellRestAPI.get_cost().""" + + def test_get_cost_single_workflow(self) -> None: + cra = CromwellRestAPI() + with ( + patch.object( + cra, + 'find_valid_workflow_ids', + return_value=['f9c26f2e-f550-4748-a650-5d0d4cab9f3a'], + ), + patch.object(cra, '_CromwellRestAPI__request_get', return_value=SAMPLE_COST_RESPONSE), + ): + result = cra.get_cost(workflow_ids=['f9c26f2e-f550-4748-a650-5d0d4cab9f3a']) + + assert result is not None + assert len(result) == 1 + assert result[0]['cost'] == 1.23 + assert result[0]['currency'] == 'USD' + + def test_get_cost_multiple_workflows(self) -> None: + cra = CromwellRestAPI() + ids = [ + 'f9c26f2e-f550-4748-a650-5d0d4cab9f3a', + 'a1b2c3d4-e5f6-7890-abcd-ef1234567890', + ] + with ( + patch.object(cra, 'find_valid_workflow_ids', return_value=ids), + patch.object( + cra, + '_CromwellRestAPI__request_get', + side_effect=[SAMPLE_COST_RESPONSE, SAMPLE_COST_RESPONSE_WITH_ERRORS], + ), + ): + result = cra.get_cost(workflow_ids=ids) + + assert result is not None + assert len(result) == 2 + + def test_get_cost_no_valid_ids(self) -> None: + cra = CromwellRestAPI() + with patch.object(cra, 'find_valid_workflow_ids', return_value=None): + result = cra.get_cost(workflow_ids=['invalid']) + + assert result is None + + def test_get_cost_skips_none_responses(self) -> None: + cra = CromwellRestAPI() + with ( + patch.object( + cra, + 'find_valid_workflow_ids', + return_value=['f9c26f2e-f550-4748-a650-5d0d4cab9f3a'], + ), + patch.object(cra, '_CromwellRestAPI__request_get', return_value=None), + ): + result = cra.get_cost(workflow_ids=['f9c26f2e-f550-4748-a650-5d0d4cab9f3a']) + + assert result is not None + assert len(result) == 0 + + +class TestCaperClientCost: + """Tests for CaperClient.cost().""" + + def test_cost_delegates_to_rest_api(self) -> None: + with patch.object(CaperClient, '__init__', lambda self, **kw: None): + client = CaperClient() + client._cromwell_rest_api = MagicMock() + client._cromwell_rest_api.get_cost.return_value = [SAMPLE_COST_RESPONSE] + + result = client.cost(wf_ids_or_labels=['f9c26f2e-f550-4748-a650-5d0d4cab9f3a']) + + assert result == [SAMPLE_COST_RESPONSE] + client._cromwell_rest_api.get_cost.assert_called_once() + + +class TestSubcmdCostFormatting: + """Tests for subcmd_cost() output formatting.""" + + def _make_args(self, cost_format: str = 'human') -> MagicMock: + args = MagicMock() + args.cost_format = cost_format + args.wf_id_or_label = ['f9c26f2e-f550-4748-a650-5d0d4cab9f3a'] + return args + + def _make_client(self, cost_results: list[dict]) -> MagicMock: + client = MagicMock() + client.cost.return_value = cost_results + return client + + def test_human_format_single_workflow(self, capsys: pytest.CaptureFixture) -> None: + client = self._make_client([SAMPLE_COST_RESPONSE]) + args = self._make_args('human') + + subcmd_cost(client, args) + + out = capsys.readouterr().out + assert 'f9c26f2e-f550-4748-a650-5d0d4cab9f3a' in out + assert '$1.23' in out + assert 'Succeeded' in out + # Total line only shown for multiple workflows or when warnings present + assert 'Total' not in out + + def test_human_format_with_warnings(self, capsys: pytest.CaptureFixture) -> None: + client = self._make_client([SAMPLE_COST_RESPONSE, SAMPLE_COST_RESPONSE_WITH_ERRORS]) + args = self._make_args('human') + + subcmd_cost(client, args) + + out = capsys.readouterr().out + assert 'Warnings' in out + assert 'incomplete' in out.lower() + assert 'vmCostPerHour' in out + + def test_tsv_format(self, capsys: pytest.CaptureFixture) -> None: + client = self._make_client([SAMPLE_COST_RESPONSE]) + args = self._make_args('tsv') + + subcmd_cost(client, args) + + out = capsys.readouterr().out + lines = out.strip().split('\n') + assert lines[0] == 'workflow_id\tstatus\tcost\tcurrency\terrors' + fields = lines[1].split('\t') + assert fields[0] == 'f9c26f2e-f550-4748-a650-5d0d4cab9f3a' + assert fields[2] == '1.23' + + def test_json_format(self, capsys: pytest.CaptureFixture) -> None: + client = self._make_client([SAMPLE_COST_RESPONSE]) + args = self._make_args('json') + + subcmd_cost(client, args) + + out = capsys.readouterr().out + parsed = json.loads(out) + assert len(parsed) == 1 + assert parsed[0]['cost'] == 1.23 + + def test_no_results(self) -> None: + client = self._make_client(None) + args = self._make_args('human') + + with pytest.raises(ValueError, match='no workflow'): + subcmd_cost(client, args) diff --git a/tests/test_cromwell_backend.py b/tests/test_cromwell_backend.py index dd8e63c..5a164e1 100644 --- a/tests/test_cromwell_backend.py +++ b/tests/test_cromwell_backend.py @@ -5,7 +5,7 @@ """ -from caper.cromwell_backend import CromwellBackendBase +from caper.cromwell_backend import CromwellBackendBase, CromwellBackendCostService def test_cromwell_backend_base_backend() -> None: @@ -156,4 +156,26 @@ def test_cromwell_backend_gcp_default_has_dockerhub_mirror() -> None: config = gcp.backend_config assert 'docker-mirror' in config assert config['docker-mirror']['dockerhub']['enabled'] is True - assert config['docker-mirror']['dockerhub']['address'] == 'mirror.gcr.io' + + +def test_cost_service_enabled() -> None: + """When enabled, generates services.GcpCostCatalogService HOCON block.""" + cs = CromwellBackendCostService(enabled=True) + assert 'services' in cs + svc = cs['services']['GcpCostCatalogService'] + assert svc['class'] == 'cromwell.services.cost.GcpCostCatalogService' + assert svc['config']['enabled'] is True + assert svc['config']['catalogExpirySeconds'] == 86400 + + +def test_cost_service_disabled() -> None: + """When disabled, generates empty dict.""" + cs = CromwellBackendCostService(enabled=False) + assert dict(cs) == {} + + +def test_cost_service_custom_expiry() -> None: + """Custom catalog expiry seconds are passed through.""" + cs = CromwellBackendCostService(enabled=True, catalog_expiry_seconds=3600) + svc = cs['services']['GcpCostCatalogService'] + assert svc['config']['catalogExpirySeconds'] == 3600