feat(job): expose estimated and actual cost on jobs

docs/cost_estimates.md

@@ -0,0 +1,48 @@
+# Cost Estimates
+
+Before you commit to a large labeling run, you can ask Rapidata what a job is expected to cost. Every job definition and every running job exposes an `estimated_cost` property that returns a `CostEstimate`.
+
+## Getting an Estimate
+
+You can read the estimate straight from a job definition, before assigning it to an audience — useful for checking the cost of a run before you launch it:
+
+```py
+from rapidata import RapidataClient
+
+client = RapidataClient()
+
+job_definition = client.job.create_compare_job_definition(
+    name="Example Image Prompt Alignment",
+    instruction="Which image matches the description better?",
+    datapoints=[
+        ["https://assets.rapidata.ai/midjourney-5.2_37_3.jpg",
+         "https://assets.rapidata.ai/flux-1-pro_37_0.jpg"]
+    ],
+    contexts=["A small blue book sitting on a large red book."],
+)
+
+estimate = job_definition.estimated_cost # (1)!
+print(f"About {estimate.estimated_cost} for {estimate.required_responses} responses")
+```
+
+1. The estimate is priced shortly after the definition is created. Reading it right away blocks for a moment until it is ready, then returns.
+
+The same property is available on a job once it is running:
+
+```py
+job = audience.assign_job(job_definition)
+print(job.estimated_cost.estimated_cost)
+```
+
+## What the Estimate Contains
+
+A `CostEstimate` has three fields:
+
+| Field | Description |
+|---|---|
+| `estimated_cost` | The estimated total cost of running the job to completion, in your account's billing currency. |
+| `datapoint_count` | The number of datapoints the job will label. |
+| `required_responses` | The total number of responses the job collects to complete. |
+
+!!! note
+    This is an **estimate, not the final bill**. It is based on a sample of the job's tasks scaled to the number of responses requested, so the amount you are actually charged can differ. Features like [Early Stopping](confidence_stopping.md) can also lower the final cost by collecting fewer responses than the maximum.

mkdocs.yml

@@ -65,6 +65,7 @@ plugins:
           - understanding_the_results.md
           - error_handling.md
           - confidence_stopping.md
+          - cost_estimates.md
           - human_prompting.md
           - config.md
         Examples:
@@ -128,6 +129,7 @@ nav:
       - Parameter Reference: job_definition_parameters.md
       - Understanding Results: understanding_the_results.md
       - Early Stopping: confidence_stopping.md
+      - Cost Estimates: cost_estimates.md
       - Instruction Design: human_prompting.md
       - Error Handling: error_handling.md
       - Logging & Config: config.md

src/rapidata/__init__.py

@@ -11,6 +11,7 @@
     RapidataJob,
     RapidataJobDefinition,
     RapidataJobManager,
+    CostEstimate,
     RapidataSignal,
     RapidataSignalManager,
     ValidationSetManager,

src/rapidata/rapidata_client/__init__.py

@@ -6,7 +6,12 @@
     RapidataFilteredAudience,
 )
 from .order import RapidataOrderManager, RapidataOrder
-from .job import RapidataJob, RapidataJobDefinition, RapidataJobManager
+from .job import (
+    RapidataJob,
+    RapidataJobDefinition,
+    RapidataJobManager,
+    CostEstimate,
+)
 from .signal import RapidataSignal, RapidataSignalManager
 from .validation import ValidationSetManager, RapidataValidationSet, Box
 from .results import RapidataResults

src/rapidata/rapidata_client/job/__init__.py

@@ -1,3 +1,4 @@
 from .rapidata_job_definition import RapidataJobDefinition
 from .rapidata_job_manager import RapidataJobManager
 from .rapidata_job import RapidataJob
+from .cost import CostEstimate

src/rapidata/rapidata_client/job/cost.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from time import monotonic, sleep
+from typing import Callable, TypeVar, TYPE_CHECKING
+
+from rapidata.rapidata_client.api.rapidata_api_client import (
+    suppress_rapidata_error_logging,
+)
+from rapidata.rapidata_client.config import logger
+from rapidata.rapidata_client.exceptions.rapidata_error import RapidataError
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.get_job_cost_estimate_endpoint_output import (
+        GetJobCostEstimateEndpointOutput,
+    )
+    from rapidata.api_client.models.get_job_definition_cost_estimate_endpoint_output import (
+        GetJobDefinitionCostEstimateEndpointOutput,
+    )
+
+T = TypeVar("T")
+
+# The estimate is not priced instantly after a job is created; the endpoint
+# answers 409 until it becomes available.
+_ESTIMATE_NOT_READY_STATUS = 409
+DEFAULT_ESTIMATE_TIMEOUT = 300.0
+DEFAULT_ESTIMATE_POLL_INTERVAL = 5.0
+
+
+@dataclass(frozen=True)
+class CostEstimate:
+    """An approximate estimate of what a job will cost to run to completion.
+
+    This is an estimate, not the final bill: it is based on a sample of the
+    job's tasks and scaled to the total number of responses requested, so the
+    amount you are actually charged can differ.
+
+    Attributes:
+        estimated_cost: The estimated total cost of running the job to completion.
+        datapoint_count: The number of datapoints the job will label.
+        required_responses: The total number of responses the job collects to complete.
+    """
+
+    estimated_cost: float
+    datapoint_count: int
+    required_responses: int
+
+    @classmethod
+    def _from_model(
+        cls,
+        model: (
+            GetJobCostEstimateEndpointOutput
+            | GetJobDefinitionCostEstimateEndpointOutput
+        ),
+    ) -> CostEstimate:
+        return cls(
+            estimated_cost=model.estimated_cost,
+            datapoint_count=model.datapoint_count,
+            required_responses=model.required_responses,
+        )
+
+
+def _poll_for_cost_estimate(
+    fetch: Callable[[], T | None],
+    *,
+    timeout: float,
+    interval: float,
+) -> T:
+    """Call ``fetch`` until the estimate is available, retrying while it is not.
+
+    Until the estimate has been priced the endpoint signals "not ready" in two
+    ways: an HTTP 409, or a success with an empty body (which the generated
+    client returns as ``None``). Both are retried until a result is available;
+    any other error is raised immediately.
+
+    Raises:
+        TimeoutError: If the estimate is still not available after ``timeout`` seconds.
+    """
+    deadline = monotonic() + timeout
+    while True:
+        try:
+            with suppress_rapidata_error_logging():
+                result = fetch()
+        except RapidataError as e:
+            if e.status_code != _ESTIMATE_NOT_READY_STATUS:
+                raise
+            result = None
+
+        if result is not None:
+            return result
+
+        if monotonic() >= deadline:
+            raise TimeoutError(
+                f"Cost estimate was not available after {timeout:.0f}s - "
+                "try again shortly."
+            )
+        logger.debug("Cost estimate not ready yet, polling...")
+        sleep(interval)

src/rapidata/rapidata_client/job/rapidata_job.py

@@ -19,6 +19,12 @@
 from rapidata.rapidata_client.api.rapidata_api_client import (
     suppress_rapidata_error_logging,
 )
+from rapidata.rapidata_client.job.cost import (
+    CostEstimate,
+    DEFAULT_ESTIMATE_POLL_INTERVAL,
+    DEFAULT_ESTIMATE_TIMEOUT,
+    _poll_for_cost_estimate,
+)
 
 if TYPE_CHECKING:
     from rapidata.rapidata_client.results.rapidata_results import RapidataResults
@@ -60,6 +66,7 @@ def __init__(
         self.definition_id = definition_id
         self.__pipeline_id = pipeline_id
         self.__completed_at = None
+        self.__estimated_cost: CostEstimate | None = None
         self.job_details_page = f"https://app.{self._openapi_service.environment}/audiences/{self.audience_id}/job/{self.id}"
         logger.debug("RapidataJob initialized")
 
@@ -152,6 +159,29 @@ def pipeline_id(self) -> str:
             ).pipeline_id
         return self.__pipeline_id
 
+    @property
+    def estimated_cost(self) -> CostEstimate:
+        """An approximate estimate of what this job will cost to run to completion.
+
+        This is an estimate, not the final bill - see :class:`CostEstimate`. The
+        estimate is priced shortly after the job is created; this call waits for
+        it to become available.
+
+        Raises:
+            TimeoutError: If the estimate is still not available after a few minutes.
+        """
+        if self.__estimated_cost is None:
+            with tracer.start_as_current_span("RapidataJob.estimated_cost"):
+                model = _poll_for_cost_estimate(
+                    lambda: self._openapi_service.order.job_api.job_job_id_cost_estimate_get(
+                        self.id
+                    ),
+                    timeout=DEFAULT_ESTIMATE_TIMEOUT,
+                    interval=DEFAULT_ESTIMATE_POLL_INTERVAL,
+                )
+                self.__estimated_cost = CostEstimate._from_model(model)
+        return self.__estimated_cost
+
     def get_status(self) -> str:
         """
         Gets the status of the job.

src/rapidata/rapidata_client/job/rapidata_job_definition.py

@@ -6,6 +6,12 @@
 from typing import Literal
 from rapidata.service.openapi_service import OpenAPIService
 from rapidata.rapidata_client.config import tracer, logger, managed_print
+from rapidata.rapidata_client.job.cost import (
+    CostEstimate,
+    DEFAULT_ESTIMATE_POLL_INTERVAL,
+    DEFAULT_ESTIMATE_TIMEOUT,
+    _poll_for_cost_estimate,
+)
 import webbrowser
 import urllib.parse
 from colorama import Fore
@@ -24,6 +30,30 @@ def __init__(
         self._job_details_page = (
             f"https://app.{self._openapi_service.environment}/definitions/{self.id}"
         )
+        self.__estimated_cost: CostEstimate | None = None
+
+    @property
+    def estimated_cost(self) -> CostEstimate:
+        """An approximate estimate of what this job will cost to run to completion.
+
+        This is an estimate, not the final bill - see :class:`CostEstimate`. The
+        estimate is priced shortly after the definition is created; this call
+        waits for it to become available.
+
+        Raises:
+            TimeoutError: If the estimate is still not available after a few minutes.
+        """
+        if self.__estimated_cost is None:
+            with tracer.start_as_current_span("RapidataJobDefinition.estimated_cost"):
+                model = _poll_for_cost_estimate(
+                    lambda: self._openapi_service.order.job_api.job_definition_definition_id_cost_estimate_get(
+                        self.id
+                    ),
+                    timeout=DEFAULT_ESTIMATE_TIMEOUT,
+                    interval=DEFAULT_ESTIMATE_POLL_INTERVAL,
+                )
+                self.__estimated_cost = CostEstimate._from_model(model)
+        return self.__estimated_cost
 
     def preview(self) -> RapidataJobDefinition:
         """Will open the browser where you can preview the job definition before giving it to an audience."""