edsl 0.1.56__py3-none-any.whl → 0.1.58__py3-none-any.whl

edsl/jobs/jobs_pricing_estimation.py

@@ -13,16 +13,15 @@ if TYPE_CHECKING:
     from ..invigilators.invigilator_base import Invigilator
 
 from .fetch_invigilator import FetchInvigilator
+from ..coop.utils import CostConverter
 from ..caching import CacheEntry
 from ..dataset import Dataset
+from ..language_models.price_manager import PriceRetriever
 
 logger = logging.getLogger(__name__)
 
 
 class PromptCostEstimator:
-
-    DEFAULT_INPUT_PRICE_PER_MILLION_TOKENS = 1.0
-    DEFAULT_OUTPUT_PRICE_PER_MILLION_TOKENS = 1.0
     CHARS_PER_TOKEN = 4
     OUTPUT_TOKENS_PER_INPUT_TOKEN = 0.75
     PIPING_MULTIPLIER = 2
@@ -37,7 +36,7 @@ class PromptCostEstimator:
     ):
         self.system_prompt = system_prompt
         self.user_prompt = user_prompt
-        self.price_lookup = price_lookup
+        self.price_retriever = PriceRetriever(price_lookup)
         self.inference_service = inference_service
         self.model = model
 
@@ -49,91 +48,6 @@ class PromptCostEstimator:
             return PromptCostEstimator.PIPING_MULTIPLIER
         return 1
 
-    def _get_fallback_price(self, inference_service: str) -> Dict:
-        """
-        Get fallback prices for a service.
-        - First fallback: The highest input and output prices for that service from the price lookup.
-        - Second fallback: $1.00 per million tokens (for both input and output).
-
-        Args:
-            inference_service (str): The inference service name
-
-        Returns:
-            Dict: Price information
-        """
-        PriceEntry = namedtuple("PriceEntry", ["tokens_per_usd", "price_info"])
-
-        service_prices = [
-            prices
-            for (service, _), prices in self.price_lookup.items()
-            if service == inference_service
-        ]
-
-        default_input_price_info = {
-            "one_usd_buys": 1_000_000,
-            "service_stated_token_qty": 1_000_000,
-            "service_stated_token_price": self.DEFAULT_INPUT_PRICE_PER_MILLION_TOKENS,
-        }
-        default_output_price_info = {
-            "one_usd_buys": 1_000_000,
-            "service_stated_token_qty": 1_000_000,
-            "service_stated_token_price": self.DEFAULT_OUTPUT_PRICE_PER_MILLION_TOKENS,
-        }
-
-        # Find the most expensive price entries (lowest tokens per USD)
-        input_price_info = default_input_price_info
-        output_price_info = default_output_price_info
-
-        input_prices = [
-            PriceEntry(float(p["input"]["one_usd_buys"]), p["input"])
-            for p in service_prices
-            if "input" in p
-        ]
-        if input_prices:
-            input_price_info = min(
-                input_prices, key=lambda price: price.tokens_per_usd
-            ).price_info
-
-        output_prices = [
-            PriceEntry(float(p["output"]["one_usd_buys"]), p["output"])
-            for p in service_prices
-            if "output" in p
-        ]
-        if output_prices:
-            output_price_info = min(
-                output_prices, key=lambda price: price.tokens_per_usd
-            ).price_info
-
-        return {
-            "input": input_price_info,
-            "output": output_price_info,
-        }
-
-    def get_price(self, inference_service: str, model: str) -> Dict:
-        """Get the price information for a specific service and model."""
-        key = (inference_service, model)
-        return self.price_lookup.get(key) or self._get_fallback_price(inference_service)
-
-    def get_price_per_million_tokens(
-        self,
-        relevant_prices: Dict,
-        token_type: Literal["input", "output"],
-    ) -> Dict:
-        """
-        Get the price per million tokens for a specific service, model, and token type.
-        """
-        service_price = relevant_prices[token_type]["service_stated_token_price"]
-        service_qty = relevant_prices[token_type]["service_stated_token_qty"]
-
-        if service_qty == 1_000_000:
-            price_per_million_tokens = service_price
-        elif service_qty == 1_000:
-            price_per_million_tokens = service_price * 1_000
-        else:
-            price_per_token = service_price / service_qty
-            price_per_million_tokens = round(price_per_token * 1_000_000, 10)
-        return price_per_million_tokens
-
     def __call__(self):
         user_prompt_chars = len(str(self.user_prompt)) * self.get_piping_multiplier(
             str(self.user_prompt)
@@ -145,13 +59,15 @@ class PromptCostEstimator:
         input_tokens = (user_prompt_chars + system_prompt_chars) // self.CHARS_PER_TOKEN
         output_tokens = math.ceil(self.OUTPUT_TOKENS_PER_INPUT_TOKEN * input_tokens)
 
-        relevant_prices = self.get_price(self.inference_service, self.model)
+        relevant_prices = self.price_retriever.get_price(
+            self.inference_service, self.model
+        )
 
-        input_price_per_million_tokens = self.get_price_per_million_tokens(
-            relevant_prices, "input"
+        input_price_per_million_tokens = (
+            self.price_retriever.get_price_per_million_tokens(relevant_prices, "input")
         )
-        output_price_per_million_tokens = self.get_price_per_million_tokens(
-            relevant_prices, "output"
+        output_price_per_million_tokens = (
+            self.price_retriever.get_price_per_million_tokens(relevant_prices, "output")
         )
 
         input_price_per_token = input_price_per_million_tokens / 1_000_000
@@ -429,8 +345,14 @@ class JobsPrompts:
             group["cost_usd"] *= iterations
             detailed_costs.append(group)
 
+        # Convert to credits
+        converter = CostConverter()
+        for group in detailed_costs:
+            group["credits_hold"] = converter.usd_to_credits(group["cost_usd"])
+
         # Calculate totals
-        estimated_total_cost = sum(group["cost_usd"] for group in detailed_costs)
+        estimated_total_cost_usd = sum(group["cost_usd"] for group in detailed_costs)
+        total_credits_hold = sum(group["credits_hold"] for group in detailed_costs)
         estimated_total_input_tokens = sum(
             group["tokens"]
             for group in detailed_costs
@@ -443,7 +365,8 @@ class JobsPrompts:
         )
 
         output = {
-            "estimated_total_cost_usd": estimated_total_cost,
+            "estimated_total_cost_usd": estimated_total_cost_usd,
+            "total_credits_hold": total_credits_hold,
             "estimated_total_input_tokens": estimated_total_input_tokens,
             "estimated_total_output_tokens": estimated_total_output_tokens,
             "detailed_costs": detailed_costs,

edsl/jobs/jobs_remote_inference_logger.py

@@ -23,15 +23,38 @@ class LogMessage:
     status: JobsStatus
 
 
+@dataclass
+class JobRunExceptionCounter:
+    exception_type: str = None
+    inference_service: str = None
+    model: str = None
+    question_name: str = None
+    exception_count: int = None
+
+
+@dataclass
+class ModelCost:
+    service: str = None
+    model: str = None
+    input_tokens: int = None
+    input_cost_usd: float = None
+    output_tokens: int = None
+    output_cost_usd: float = None
+
+
 @dataclass
 class JobsInfo:
     job_uuid: str = None
     progress_bar_url: str = None
     error_report_url: str = None
+    remote_inference_url: str = None
+    remote_cache_url: str = None
     results_uuid: str = None
     results_url: str = None
     completed_interviews: int = None
     failed_interviews: int = None
+    exception_summary: list[JobRunExceptionCounter] = None
+    model_costs: list[ModelCost] = None
 
     pretty_names = {
         "job_uuid": "Job UUID",
@@ -39,6 +62,8 @@ class JobsInfo:
         "error_report_url": "Exceptions Report URL",
         "results_uuid": "Results UUID",
         "results_url": "Results URL",
+        "remote_inference_url": "Remote Jobs",
+        "remote_cache_url": "Remote Cache",
     }
 
 
@@ -57,6 +82,8 @@ class JobLogger(ABC):
             "results_url",
             "completed_interviews",
             "failed_interviews",
+            "model_costs",
+            "exception_summary",
         ],
         value: str,
     ):

edsl/jobs/jobs_runner_status.py

@@ -11,15 +11,7 @@ from uuid import UUID
 
 if TYPE_CHECKING:
     from .jobs import Jobs
-
-
-@dataclass
-class ModelInfo:
-    model_name: str
-    TPM_limit_k: float
-    RPM_limit_k: float
-    num_tasks_waiting: int
-    token_usage_info: dict
+    from ..interviews import Interview
 
 
 class StatisticsTracker:
@@ -29,16 +21,33 @@ class StatisticsTracker:
         self.completed_count = 0
         self.completed_by_model = defaultdict(int)
         self.distinct_models = distinct_models
+        self.interviews_with_exceptions = 0
         self.total_exceptions = 0
         self.unfixed_exceptions = 0
+        self.exceptions_counter = defaultdict(int)
 
     def add_completed_interview(
-        self, model: str, num_exceptions: int = 0, num_unfixed: int = 0
+        self,
+        model: str,
+        exceptions: list[dict],
+        num_exceptions: int = 0,
+        num_unfixed: int = 0,
     ):
         self.completed_count += 1
         self.completed_by_model[model] += 1
         self.total_exceptions += num_exceptions
         self.unfixed_exceptions += num_unfixed
+        if num_exceptions > 0:
+            self.interviews_with_exceptions += 1
+
+        for exception in exceptions:
+            key = (
+                exception["exception_type"],
+                exception["inference_service"],
+                exception["model"],
+                exception["question_name"],
+            )
+            self.exceptions_counter[key] += 1
 
     def get_elapsed_time(self) -> float:
         return time.time() - self.start_time
@@ -88,9 +97,7 @@ class JobsRunnerStatusBase(ABC):
         ]
         self.num_total_interviews = n * len(self.jobs)
 
-        self.distinct_models = list(
-            set(model.model for model in self.jobs.models)
-        )
+        self.distinct_models = list(set(model.model for model in self.jobs.models))
 
         self.stats_tracker = StatisticsTracker(
             total_interviews=self.num_total_interviews,
@@ -130,6 +137,7 @@ class JobsRunnerStatusBase(ABC):
         status_dict = {
             "overall_progress": {
                 "completed": self.stats_tracker.completed_count,
+                "has_exceptions": self.stats_tracker.interviews_with_exceptions,
                 "total": self.num_total_interviews,
                 "percent": (
                     (
@@ -148,16 +156,36 @@ class JobsRunnerStatusBase(ABC):
                 if self.stats_tracker.completed_count >= self.num_total_interviews
                 else "running"
             ),
+            "exceptions_counter": [
+                {
+                    "exception_type": exception_type,
+                    "inference_service": inference_service,
+                    "model": model,
+                    "question_name": question_name,
+                    "count": count,
+                }
+                for (
+                    exception_type,
+                    inference_service,
+                    model,
+                    question_name,
+                ), count in self.stats_tracker.exceptions_counter.items()
+            ],
         }
 
         model_queues = {}
         # Check if bucket collection exists and is not empty
-        if (hasattr(self.jobs, 'run_config') and 
-            hasattr(self.jobs.run_config, 'environment') and 
-            hasattr(self.jobs.run_config.environment, 'bucket_collection') and 
-            self.jobs.run_config.environment.bucket_collection):
-            
-            for model, bucket in self.jobs.run_config.environment.bucket_collection.items():
+        if (
+            hasattr(self.jobs, "run_config")
+            and hasattr(self.jobs.run_config, "environment")
+            and hasattr(self.jobs.run_config.environment, "bucket_collection")
+            and self.jobs.run_config.environment.bucket_collection
+        ):
+
+            for (
+                model,
+                bucket,
+            ) in self.jobs.run_config.environment.bucket_collection.items():
                 model_name = model.model
                 model_queues[model_name] = {
                     "language_model_name": model_name,
@@ -166,7 +194,9 @@ class JobsRunnerStatusBase(ABC):
                         "requested": bucket.requests_bucket.num_requests,
                         "tokens_returned": bucket.requests_bucket.tokens_returned,
                         "target_rate": round(bucket.requests_bucket.target_rate, 1),
-                        "current_rate": round(bucket.requests_bucket.get_throughput(), 1),
+                        "current_rate": round(
+                            bucket.requests_bucket.get_throughput(), 1
+                        ),
                     },
                     "tokens_bucket": {
                         "completed": bucket.tokens_bucket.num_released,
@@ -179,10 +209,11 @@ class JobsRunnerStatusBase(ABC):
         status_dict["language_model_queues"] = model_queues
         return status_dict
 
-    def add_completed_interview(self, interview):
+    def add_completed_interview(self, interview: "Interview"):
         """Records a completed interview without storing the full interview data."""
         self.stats_tracker.add_completed_interview(
             model=interview.model.model,
+            exceptions=interview.exceptions.list(),
             num_exceptions=interview.exceptions.num_exceptions(),
             num_unfixed=interview.exceptions.num_unfixed_exceptions(),
         )

edsl/jobs/remote_inference.py

@@ -1,11 +1,12 @@
 import re
+import math
 from typing import Optional, Union, Literal, TYPE_CHECKING, NewType, Callable, Any
 from dataclasses import dataclass
 from ..coop import CoopServerResponseError
-from ..coop.utils import VisibilityType
+from ..coop.utils import VisibilityType, CostConverter
 from ..coop.coop import RemoteInferenceResponse, RemoteInferenceCreationInfo
 from .jobs_status_enums import JobsStatus
-from .jobs_remote_inference_logger import JobLogger
+from .jobs_remote_inference_logger import JobLogger, JobRunExceptionCounter, ModelCost
 from .exceptions import RemoteInferenceError
 
 
@@ -94,6 +95,12 @@ class JobsRemoteInferenceHandler:
             "Remote inference activated. Sending job to server...",
             status=JobsStatus.QUEUED,
         )
+        logger.add_info(
+            "remote_inference_url", f"{self.expected_parrot_url}/home/remote-inference"
+        )
+        logger.add_info(
+            "remote_cache_url", f"{self.expected_parrot_url}/home/remote-cache"
+        )
         remote_job_creation_data = coop.remote_inference_create(
             self.jobs,
             description=remote_inference_description,
@@ -183,7 +190,9 @@ class JobsRemoteInferenceHandler:
         self, job_info: RemoteJobInfo, remote_job_data: RemoteInferenceResponse
     ) -> None:
         "Handles a failed job by logging the error and updating the job status."
-        latest_error_report_url = remote_job_data.get("latest_error_report_url")
+        error_report_url = remote_job_data.get("latest_job_run_details", {}).get(
+            "error_report_url"
+        )
 
         reason = remote_job_data.get("reason")
 
@@ -193,8 +202,8 @@ class JobsRemoteInferenceHandler:
                 status=JobsStatus.FAILED,
             )
 
-        if latest_error_report_url:
-            job_info.logger.add_info("error_report_url", latest_error_report_url)
+        if error_report_url:
+            job_info.logger.add_info("error_report_url", error_report_url)
 
         job_info.logger.update("Job failed.", status=JobsStatus.FAILED)
         job_info.logger.update(
@@ -206,39 +215,46 @@ class JobsRemoteInferenceHandler:
             status=JobsStatus.FAILED,
         )
 
-    def _handle_partially_failed_job_interview_details(
+    def _update_interview_details(
         self, job_info: RemoteJobInfo, remote_job_data: RemoteInferenceResponse
     ) -> None:
-        "Extracts the interview details from the remote job data."
-        try:
-            # Job details is a string of the form "64 out of 1,758 interviews failed"
-            job_details = remote_job_data.get("latest_failure_description")
-
-            text_without_commas = job_details.replace(",", "")
-
-            # Find all numbers in the text
-            numbers = [int(num) for num in re.findall(r"\d+", text_without_commas)]
-
-            failed = numbers[0]
-            total = numbers[1]
-            completed = total - failed
-
-            job_info.logger.add_info("completed_interviews", completed)
-            job_info.logger.add_info("failed_interviews", failed)
-        # This is mainly helpful metadata, and any errors here should not stop the code
-        except:
-            pass
+        "Updates the interview details in the job info."
+        latest_job_run_details = remote_job_data.get("latest_job_run_details", {})
+        interview_details = latest_job_run_details.get("interview_details", {}) or {}
+        completed_interviews = interview_details.get("completed_interviews", 0)
+        interviews_with_exceptions = interview_details.get(
+            "interviews_with_exceptions", 0
+        )
+        interviews_without_exceptions = (
+            completed_interviews - interviews_with_exceptions
+        )
+        job_info.logger.add_info("completed_interviews", interviews_without_exceptions)
+        job_info.logger.add_info("failed_interviews", interviews_with_exceptions)
+
+        exception_summary = interview_details.get("exception_summary", []) or []
+        if exception_summary:
+            job_run_exception_counters = []
+            for exception in exception_summary:
+                exception_counter = JobRunExceptionCounter(
+                    exception_type=exception.get("exception_type"),
+                    inference_service=exception.get("inference_service"),
+                    model=exception.get("model"),
+                    question_name=exception.get("question_name"),
+                    exception_count=exception.get("exception_count"),
+                )
+                job_run_exception_counters.append(exception_counter)
+            job_info.logger.add_info("exception_summary", job_run_exception_counters)
 
     def _handle_partially_failed_job(
         self, job_info: RemoteJobInfo, remote_job_data: RemoteInferenceResponse
     ) -> None:
         "Handles a partially failed job by logging the error and updating the job status."
-        self._handle_partially_failed_job_interview_details(job_info, remote_job_data)
-
-        latest_error_report_url = remote_job_data.get("latest_error_report_url")
+        error_report_url = remote_job_data.get("latest_job_run_details", {}).get(
+            "error_report_url"
+        )
 
-        if latest_error_report_url:
-            job_info.logger.add_info("error_report_url", latest_error_report_url)
+        if error_report_url:
+            job_info.logger.add_info("error_report_url", error_report_url)
 
         job_info.logger.update(
             "Job completed with partial results.", status=JobsStatus.PARTIALLY_FAILED
@@ -263,6 +279,128 @@ class JobsRemoteInferenceHandler:
         )
         time.sleep(self.poll_interval)
 
+    def _get_expenses_from_results(
+        self, results: "Results", include_cached_responses_in_cost: bool = False
+    ) -> dict:
+        """
+        Calculates expenses from Results object.
+
+        Args:
+            results: Results object containing model responses
+            include_cached_responses_in_cost: Whether to include cached responses in cost calculation
+
+        Returns:
+            Dictionary mapping ExpenseKey to TokenExpense information
+        """
+        expenses = {}
+
+        for result in results:
+            raw_response = result["raw_model_response"]
+
+            # Process each cost field in the response
+            for key in raw_response:
+                if not key.endswith("_cost"):
+                    continue
+
+                result_cost = raw_response[key]
+                if not isinstance(result_cost, (int, float)):
+                    continue
+
+                question_name = key.removesuffix("_cost")
+                cache_used = result["cache_used_dict"][question_name]
+
+                # Skip if we're excluding cached responses and this was cached
+                if not include_cached_responses_in_cost and cache_used:
+                    continue
+
+                # Get expense keys for input and output tokens
+                input_key = (
+                    result["model"]._inference_service_,
+                    result["model"].model,
+                    "input",
+                    raw_response[f"{question_name}_input_price_per_million_tokens"],
+                )
+                output_key = (
+                    result["model"]._inference_service_,
+                    result["model"].model,
+                    "output",
+                    raw_response[f"{question_name}_output_price_per_million_tokens"],
+                )
+
+                # Update input token expenses
+                if input_key not in expenses:
+                    expenses[input_key] = {
+                        "tokens": 0,
+                        "cost_usd": 0,
+                    }
+
+                input_price_per_million_tokens = input_key[3]
+                input_tokens = raw_response[f"{question_name}_input_tokens"]
+                input_cost = (input_price_per_million_tokens / 1_000_000) * input_tokens
+
+                expenses[input_key]["tokens"] += input_tokens
+                expenses[input_key]["cost_usd"] += input_cost
+
+                # Update output token expenses
+                if output_key not in expenses:
+                    expenses[output_key] = {
+                        "tokens": 0,
+                        "cost_usd": 0,
+                    }
+
+                output_price_per_million_tokens = output_key[3]
+                output_tokens = raw_response[f"{question_name}_output_tokens"]
+                output_cost = (
+                    output_price_per_million_tokens / 1_000_000
+                ) * output_tokens
+
+                expenses[output_key]["tokens"] += output_tokens
+                expenses[output_key]["cost_usd"] += output_cost
+
+        expenses_by_model = {}
+        for expense_key, expense_usage in expenses.items():
+            service, model, token_type, _ = expense_key
+            model_key = (service, model)
+
+            if model_key not in expenses_by_model:
+                expenses_by_model[model_key] = {
+                    "service": service,
+                    "model": model,
+                    "input_tokens": 0,
+                    "input_cost_usd": 0,
+                    "output_tokens": 0,
+                    "output_cost_usd": 0,
+                }
+
+            if token_type == "input":
+                expenses_by_model[model_key]["input_tokens"] += expense_usage["tokens"]
+                expenses_by_model[model_key]["input_cost_usd"] += expense_usage[
+                    "cost_usd"
+                ]
+            elif token_type == "output":
+                expenses_by_model[model_key]["output_tokens"] += expense_usage["tokens"]
+                expenses_by_model[model_key]["output_cost_usd"] += expense_usage[
+                    "cost_usd"
+                ]
+
+        converter = CostConverter()
+        for model_key, model_cost_dict in expenses_by_model.items():
+            input_cost = model_cost_dict["input_cost_usd"]
+            output_cost = model_cost_dict["output_cost_usd"]
+            model_cost_dict["input_cost_credits"] = converter.usd_to_credits(input_cost)
+            model_cost_dict["output_cost_credits"] = converter.usd_to_credits(
+                output_cost
+            )
+            # Convert back to USD (to get the rounded value)
+            model_cost_dict["input_cost_usd"] = converter.credits_to_usd(
+                model_cost_dict["input_cost_credits"]
+            )
+            model_cost_dict["output_cost_usd"] = converter.credits_to_usd(
+                model_cost_dict["output_cost_credits"]
+            )
+
+        return list(expenses_by_model.values())
+
     def _fetch_results_and_log(
         self,
         job_info: RemoteJobInfo,
@@ -274,12 +412,30 @@ class JobsRemoteInferenceHandler:
         "Fetches the results object and logs the results URL."
         job_info.logger.add_info("results_uuid", results_uuid)
         results = object_fetcher(results_uuid, expected_object_type="results")
+
+        model_cost_dicts = self._get_expenses_from_results(results)
+
+        model_costs = [
+            ModelCost(
+                service=model_cost_dict.get("service"),
+                model=model_cost_dict.get("model"),
+                input_tokens=model_cost_dict.get("input_tokens"),
+                input_cost_usd=model_cost_dict.get("input_cost_usd"),
+                output_tokens=model_cost_dict.get("output_tokens"),
+                output_cost_usd=model_cost_dict.get("output_cost_usd"),
+            )
+            for model_cost_dict in model_cost_dicts
+        ]
+        job_info.logger.add_info("model_costs", model_costs)
+
         results_url = remote_job_data.get("results_url")
         if "localhost" in results_url:
             results_url = results_url.replace("8000", "1234")
         job_info.logger.add_info("results_url", results_url)
 
         if job_status == "completed":
+            job_info.logger.add_info("completed_interviews", len(results))
+            job_info.logger.add_info("failed_interviews", 0)
             job_info.logger.update(
                 f"Job completed and Results stored on Coop. [View Results]({results_url})",
                 status=JobsStatus.COMPLETED,
@@ -302,6 +458,7 @@ class JobsRemoteInferenceHandler:
     ) -> Union[None, "Results", Literal["continue"]]:
         """Makes one attempt to fetch and process a remote job's status and results."""
         remote_job_data = remote_job_data_fetcher(job_info.job_uuid)
+        self._update_interview_details(job_info, remote_job_data)
         status = remote_job_data.get("status")
         reason = remote_job_data.get("reason")
         if status == "cancelled":

edsl/language_models/language_model.py

@@ -1019,7 +1019,7 @@ class LanguageModel(
 
         # Combine model name and parameters
         return (
-            f"Model(model_name = '{self.model}'"
+            f"Model(model_name = '{self.model}', service_name = '{self._inference_service_}'"
             + (f", {param_string}" if param_string else "")
             + ")"
         )