From 3c0c26192bfae783610c454111f60182443283a2 Mon Sep 17 00:00:00 2001 From: Dmytro Apollonin Date: Tue, 2 Jun 2026 14:02:50 -0600 Subject: [PATCH 1/4] Add ModelRun.total_cost and total_data_rows Expose a model run's total cost and processed data-row count, fetched in real time from Model Foundry (modelFoundryModelRunInfo) on property access and cached on the instance -- nothing is persisted on the run. Returns None for runs not backed by a Foundry model job. Requires the matching backend changes that surface totalDataRows on the modelFoundryModelRunInfo GraphQL field. --- .../labelbox/src/labelbox/schema/model_run.py | 48 +++++++++++++ .../tests/unit/test_unit_model_run.py | 67 +++++++++++++++++++ 2 files changed, 115 insertions(+) create mode 100644 libs/labelbox/tests/unit/test_unit_model_run.py diff --git a/libs/labelbox/src/labelbox/schema/model_run.py b/libs/labelbox/src/labelbox/schema/model_run.py index a7712858e..052fe0ca4 100644 --- a/libs/labelbox/src/labelbox/schema/model_run.py +++ b/libs/labelbox/src/labelbox/schema/model_run.py @@ -16,6 +16,8 @@ Union, ) +from lbox.exceptions import LabelboxError + from labelbox.orm.db_object import DbObject, experimental from labelbox.orm.model import Entity, Field, Relationship from labelbox.orm.query import results_query_part @@ -65,6 +67,52 @@ class Status(Enum): COMPLETE = "COMPLETE" FAILED = "FAILED" + def _get_cost_and_usage(self) -> Dict[str, Any]: + """Lazily fetches and caches cost and data row count for this Model Run. + + The data is rehydrated in real time from Model Foundry (which in turn + sources it from the model service); nothing is persisted on the Model + Run itself. Returns an empty dict for Model Runs that were not produced + by a Foundry app (i.e. that have no associated model job). + """ + if getattr(self, "_cost_and_usage", None) is None: + query_str = """ + query GetModelRunCostInfoPyApi($modelRunId: ID!) { + modelFoundryModelRunInfo(where: {modelRunId: $modelRunId}) { + cost + status + totalDataRows + } + } + """ + try: + res = self.client.execute(query_str, {"modelRunId": self.uid}) + self._cost_and_usage = res["modelFoundryModelRunInfo"] or {} + except LabelboxError: + # Model Runs not backed by a Foundry model job have no + # cost/usage info to report. + self._cost_and_usage = {} + return self._cost_and_usage + + @property + def total_cost(self) -> Optional[float]: + """Total cost (USD) of this Model Run, fetched in real time from Model + Foundry. ``None`` if the run is not Foundry-backed or cost is not yet + available. + """ + return self._get_cost_and_usage().get("cost") + + @property + def total_data_rows(self) -> Optional[int]: + """Number of data rows processed by this Model Run, fetched in real time + from Model Foundry. ``None`` if the run is not Foundry-backed. + """ + return self._get_cost_and_usage().get("totalDataRows") + + def refresh_cost_and_usage(self) -> None: + """Clears the cached cost/usage so the next access re-fetches live data.""" + self._cost_and_usage = None + def upsert_labels( self, label_ids: Optional[List[str]] = None, diff --git a/libs/labelbox/tests/unit/test_unit_model_run.py b/libs/labelbox/tests/unit/test_unit_model_run.py new file mode 100644 index 000000000..dbe488c1d --- /dev/null +++ b/libs/labelbox/tests/unit/test_unit_model_run.py @@ -0,0 +1,67 @@ +from unittest.mock import MagicMock + +from lbox.exceptions import LabelboxError + +from labelbox.schema.model_run import ModelRun + + +def _make_model_run(client): + return ModelRun( + client, + { + "id": "model-run-1", + "name": "test run", + "createdAt": "2021-06-01T00:00:00.000Z", + "updatedAt": "2021-06-01T00:00:00.000Z", + "createdBy": "user-1", + "modelId": "model-1", + "trainingMetadata": {}, + "modelAppId": "app-1", + }, + ) + + +def test_total_cost_and_data_rows_are_fetched_and_cached(): + client = MagicMock() + client.execute.return_value = { + "modelFoundryModelRunInfo": { + "cost": 3.5, + "status": "finished", + "totalDataRows": 12, + } + } + model_run = _make_model_run(client) + + assert model_run.total_cost == 3.5 + assert model_run.total_data_rows == 12 + + # Cost/usage is rehydrated once and cached across property reads. + assert client.execute.call_count == 1 + # The model run id is passed to the Foundry query. + assert client.execute.call_args[0][1] == {"modelRunId": "model-run-1"} + + +def test_refresh_cost_and_usage_refetches(): + client = MagicMock() + client.execute.return_value = { + "modelFoundryModelRunInfo": { + "cost": 1.0, + "status": "finished", + "totalDataRows": 1, + } + } + model_run = _make_model_run(client) + + assert model_run.total_cost == 1.0 + model_run.refresh_cost_and_usage() + assert model_run.total_cost == 1.0 + assert client.execute.call_count == 2 + + +def test_cost_and_usage_none_for_non_foundry_run(): + client = MagicMock() + client.execute.side_effect = LabelboxError("model job not found") + model_run = _make_model_run(client) + + assert model_run.total_cost is None + assert model_run.total_data_rows is None From ff2aaddcec03f2a49c46463d731dbabd5508ab1a Mon Sep 17 00:00:00 2001 From: Dmytro Apollonin Date: Tue, 2 Jun 2026 14:52:24 -0600 Subject: [PATCH 2/4] Only swallow not-found errors when rehydrating model run cost Catch ResourceNotFoundError/InternalServerError (run has no Foundry model job) and cache the empty result, but let transient errors (network, timeout, rate limit) propagate so they are not permanently cached as None. Addresses Bugbot review feedback. --- .../labelbox/src/labelbox/schema/model_run.py | 8 ++-- .../tests/unit/test_unit_model_run.py | 39 +++++++++++++++++-- 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/model_run.py b/libs/labelbox/src/labelbox/schema/model_run.py index 052fe0ca4..7f2ede90e 100644 --- a/libs/labelbox/src/labelbox/schema/model_run.py +++ b/libs/labelbox/src/labelbox/schema/model_run.py @@ -16,7 +16,7 @@ Union, ) -from lbox.exceptions import LabelboxError +from lbox.exceptions import InternalServerError, ResourceNotFoundError from labelbox.orm.db_object import DbObject, experimental from labelbox.orm.model import Entity, Field, Relationship @@ -88,9 +88,11 @@ def _get_cost_and_usage(self) -> Dict[str, Any]: try: res = self.client.execute(query_str, {"modelRunId": self.uid}) self._cost_and_usage = res["modelFoundryModelRunInfo"] or {} - except LabelboxError: + except (ResourceNotFoundError, InternalServerError): # Model Runs not backed by a Foundry model job have no - # cost/usage info to report. + # cost/usage info to report; cache the empty result. Transient + # errors (network, timeout, rate limit) are intentionally not + # caught so they propagate and the next access can retry. self._cost_and_usage = {} return self._cost_and_usage diff --git a/libs/labelbox/tests/unit/test_unit_model_run.py b/libs/labelbox/tests/unit/test_unit_model_run.py index dbe488c1d..f70d4402a 100644 --- a/libs/labelbox/tests/unit/test_unit_model_run.py +++ b/libs/labelbox/tests/unit/test_unit_model_run.py @@ -1,6 +1,11 @@ from unittest.mock import MagicMock -from lbox.exceptions import LabelboxError +import pytest +from lbox.exceptions import ( + InternalServerError, + NetworkError, + ResourceNotFoundError, +) from labelbox.schema.model_run import ModelRun @@ -58,10 +63,38 @@ def test_refresh_cost_and_usage_refetches(): assert client.execute.call_count == 2 -def test_cost_and_usage_none_for_non_foundry_run(): +@pytest.mark.parametrize( + "error", + [ + ResourceNotFoundError(message="model run not found"), + InternalServerError("no model job for run"), + ], +) +def test_cost_and_usage_none_for_non_foundry_run(error): client = MagicMock() - client.execute.side_effect = LabelboxError("model job not found") + client.execute.side_effect = error model_run = _make_model_run(client) assert model_run.total_cost is None assert model_run.total_data_rows is None + + +def test_transient_errors_propagate_and_are_not_cached(): + client = MagicMock() + client.execute.side_effect = NetworkError(Exception("boom")) + model_run = _make_model_run(client) + + with pytest.raises(NetworkError): + _ = model_run.total_cost + + # The failure is not cached, so a later successful access recovers. + client.execute.side_effect = None + client.execute.return_value = { + "modelFoundryModelRunInfo": { + "cost": 2.0, + "status": "finished", + "totalDataRows": 5, + } + } + assert model_run.total_cost == 2.0 + assert model_run.total_data_rows == 5 From 9f91a2b93478a5a5616047b7b30fb44a8d9ca5bd Mon Sep 17 00:00:00 2001 From: Dmytro Apollonin Date: Tue, 2 Jun 2026 15:36:09 -0600 Subject: [PATCH 3/4] Guard None return when rehydrating model run cost execute() returns None (not raises) for a RESOURCE_NOT_FOUND response since raise_return_resource_not_found defaults to False, so index the payload defensively to avoid a TypeError. Addresses Bugbot feedback. --- libs/labelbox/src/labelbox/schema/model_run.py | 9 +++++++-- libs/labelbox/tests/unit/test_unit_model_run.py | 16 ++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/model_run.py b/libs/labelbox/src/labelbox/schema/model_run.py index 7f2ede90e..be57ddc6c 100644 --- a/libs/labelbox/src/labelbox/schema/model_run.py +++ b/libs/labelbox/src/labelbox/schema/model_run.py @@ -87,13 +87,18 @@ def _get_cost_and_usage(self) -> Dict[str, Any]: """ try: res = self.client.execute(query_str, {"modelRunId": self.uid}) - self._cost_and_usage = res["modelFoundryModelRunInfo"] or {} except (ResourceNotFoundError, InternalServerError): # Model Runs not backed by a Foundry model job have no # cost/usage info to report; cache the empty result. Transient # errors (network, timeout, rate limit) are intentionally not # caught so they propagate and the next access can retry. - self._cost_and_usage = {} + res = None + # execute() returns None for a RESOURCE_NOT_FOUND response (it does + # not raise unless raise_return_resource_not_found=True), so guard + # against a missing payload before indexing into it. + self._cost_and_usage = (res or {}).get( + "modelFoundryModelRunInfo" + ) or {} return self._cost_and_usage @property diff --git a/libs/labelbox/tests/unit/test_unit_model_run.py b/libs/labelbox/tests/unit/test_unit_model_run.py index f70d4402a..c345072a0 100644 --- a/libs/labelbox/tests/unit/test_unit_model_run.py +++ b/libs/labelbox/tests/unit/test_unit_model_run.py @@ -79,6 +79,22 @@ def test_cost_and_usage_none_for_non_foundry_run(error): assert model_run.total_data_rows is None +@pytest.mark.parametrize( + "execute_result", + [ + None, # RESOURCE_NOT_FOUND -> execute() returns None, does not raise + {"modelFoundryModelRunInfo": None}, + ], +) +def test_cost_and_usage_none_when_payload_missing(execute_result): + client = MagicMock() + client.execute.return_value = execute_result + model_run = _make_model_run(client) + + assert model_run.total_cost is None + assert model_run.total_data_rows is None + + def test_transient_errors_propagate_and_are_not_cached(): client = MagicMock() client.execute.side_effect = NetworkError(Exception("boom")) From d0023b44ee203a4dd5b213b58fcbec249446abea Mon Sep 17 00:00:00 2001 From: Dmytro Apollonin Date: Mon, 8 Jun 2026 10:00:55 -0600 Subject: [PATCH 4/4] Trim internal implementation details from model run cost/usage docs --- .../labelbox/src/labelbox/schema/model_run.py | 24 +++++++------------ .../tests/unit/test_unit_model_run.py | 4 ++-- 2 files changed, 11 insertions(+), 17 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/model_run.py b/libs/labelbox/src/labelbox/schema/model_run.py index be57ddc6c..0791e15e5 100644 --- a/libs/labelbox/src/labelbox/schema/model_run.py +++ b/libs/labelbox/src/labelbox/schema/model_run.py @@ -70,10 +70,7 @@ class Status(Enum): def _get_cost_and_usage(self) -> Dict[str, Any]: """Lazily fetches and caches cost and data row count for this Model Run. - The data is rehydrated in real time from Model Foundry (which in turn - sources it from the model service); nothing is persisted on the Model - Run itself. Returns an empty dict for Model Runs that were not produced - by a Foundry app (i.e. that have no associated model job). + Returns an empty dict when no cost/usage information is available. """ if getattr(self, "_cost_and_usage", None) is None: query_str = """ @@ -88,14 +85,10 @@ def _get_cost_and_usage(self) -> Dict[str, Any]: try: res = self.client.execute(query_str, {"modelRunId": self.uid}) except (ResourceNotFoundError, InternalServerError): - # Model Runs not backed by a Foundry model job have no - # cost/usage info to report; cache the empty result. Transient - # errors (network, timeout, rate limit) are intentionally not + # No cost/usage info available; cache the empty result. + # Transient errors (network, timeout, rate limit) are not # caught so they propagate and the next access can retry. res = None - # execute() returns None for a RESOURCE_NOT_FOUND response (it does - # not raise unless raise_return_resource_not_found=True), so guard - # against a missing payload before indexing into it. self._cost_and_usage = (res or {}).get( "modelFoundryModelRunInfo" ) or {} @@ -103,16 +96,17 @@ def _get_cost_and_usage(self) -> Dict[str, Any]: @property def total_cost(self) -> Optional[float]: - """Total cost (USD) of this Model Run, fetched in real time from Model - Foundry. ``None`` if the run is not Foundry-backed or cost is not yet - available. + """Total cost (USD) of this Model Run. + + ``None`` if cost is not available for this run. """ return self._get_cost_and_usage().get("cost") @property def total_data_rows(self) -> Optional[int]: - """Number of data rows processed by this Model Run, fetched in real time - from Model Foundry. ``None`` if the run is not Foundry-backed. + """Number of data rows processed by this Model Run. + + ``None`` if not available for this run. """ return self._get_cost_and_usage().get("totalDataRows") diff --git a/libs/labelbox/tests/unit/test_unit_model_run.py b/libs/labelbox/tests/unit/test_unit_model_run.py index c345072a0..747f1aa00 100644 --- a/libs/labelbox/tests/unit/test_unit_model_run.py +++ b/libs/labelbox/tests/unit/test_unit_model_run.py @@ -42,7 +42,7 @@ def test_total_cost_and_data_rows_are_fetched_and_cached(): # Cost/usage is rehydrated once and cached across property reads. assert client.execute.call_count == 1 - # The model run id is passed to the Foundry query. + # The model run id is passed to the query. assert client.execute.call_args[0][1] == {"modelRunId": "model-run-1"} @@ -82,7 +82,7 @@ def test_cost_and_usage_none_for_non_foundry_run(error): @pytest.mark.parametrize( "execute_result", [ - None, # RESOURCE_NOT_FOUND -> execute() returns None, does not raise + None, # execute() can return None instead of a payload {"modelFoundryModelRunInfo": None}, ], )