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

edsl/jobs/jobs_pricing_estimation.py

@@ -1,7 +1,8 @@
 import logging
 import math
 
-from typing import List, TYPE_CHECKING, Union, Literal
+from typing import List, TYPE_CHECKING, Union, Literal, Dict
+from collections import namedtuple
 
 if TYPE_CHECKING:
     from .jobs import Jobs
@@ -20,8 +21,8 @@ logger = logging.getLogger(__name__)
 
 class PromptCostEstimator:
 
-    DEFAULT_INPUT_PRICE_PER_TOKEN = 0.000001
-    DEFAULT_OUTPUT_PRICE_PER_TOKEN = 0.000001
+    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
@@ -48,81 +49,90 @@ class PromptCostEstimator:
             return PromptCostEstimator.PIPING_MULTIPLIER
         return 1
 
-    @property
-    def key(self):
-        return (self.inference_service, self.model)
-
-    @property
-    def relevant_prices(self):
-        try:
-            return self.price_lookup[self.key]
-        except KeyError:
-            return {}
-
-    def _get_highest_price_for_service(self, price_type: str) -> Union[float, None]:
-        """Returns the highest price per token for a given service and price type (input/output).
+    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:
-            price_type: Either "input" or "output"
+            inference_service (str): The inference service name
 
         Returns:
-            float | None: The highest price per token for the service, or None if not found
+            Dict: Price information
         """
-        prices_for_service = [
-            prices[price_type]["service_stated_token_price"]
-            / prices[price_type]["service_stated_token_qty"]
+        PriceEntry = namedtuple("PriceEntry", ["tokens_per_usd", "price_info"])
+
+        service_prices = [
+            prices
             for (service, _), prices in self.price_lookup.items()
-            if service == self.inference_service and price_type in prices
+            if service == inference_service
         ]
-        return max(prices_for_service) if prices_for_service else None
 
-    def input_price_per_token(
-        self,
-    ) -> tuple[float, Literal["price_lookup", "highest_price_for_service", "default"]]:
-        try:
-            return (
-                self.relevant_prices["input"]["service_stated_token_price"]
-                / self.relevant_prices["input"]["service_stated_token_qty"]
-            ), "price_lookup"
-        except KeyError:
-            highest_price = self._get_highest_price_for_service("input")
-            if highest_price is not None:
-                import warnings
-
-                warnings.warn(
-                    f"Price data not found for {self.key}. Using highest available input price for {self.inference_service}: ${highest_price:.6f} per token"
-                )
-                return highest_price, "highest_price_for_service"
-            import warnings
+        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,
+        }
 
-            warnings.warn(
-                f"Price data not found for {self.inference_service}. Using default estimate for input token price: $1.00 / 1M tokens"
-            )
-            return self.DEFAULT_INPUT_PRICE_PER_TOKEN, "default"
+        # Find the most expensive price entries (lowest tokens per USD)
+        input_price_info = default_input_price_info
+        output_price_info = default_output_price_info
 
-    def output_price_per_token(
-        self,
-    ) -> tuple[float, Literal["price_lookup", "highest_price_for_service", "default"]]:
-        try:
-            return (
-                self.relevant_prices["output"]["service_stated_token_price"]
-                / self.relevant_prices["output"]["service_stated_token_qty"]
-            ), "price_lookup"
-        except KeyError:
-            highest_price = self._get_highest_price_for_service("output")
-            if highest_price is not None:
-                import warnings
-
-                warnings.warn(
-                    f"Price data not found for {self.key}. Using highest available output price for {self.inference_service}: ${highest_price:.6f} per token"
-                )
-                return highest_price, "highest_price_for_service"
-            import warnings
+        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
 
-            warnings.warn(
-                f"Price data not found for {self.inference_service}. Using default estimate for output token price: $1.00 / 1M tokens"
-            )
-            return self.DEFAULT_OUTPUT_PRICE_PER_TOKEN, "default"
+        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(
@@ -135,20 +145,28 @@ 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)
 
-        input_price_per_token, input_price_source = self.input_price_per_token()
-        output_price_per_token, output_price_source = self.output_price_per_token()
+        relevant_prices = self.get_price(self.inference_service, self.model)
 
-        cost = (
-            input_tokens * input_price_per_token
-            + output_tokens * output_price_per_token
+        input_price_per_million_tokens = self.get_price_per_million_tokens(
+            relevant_prices, "input"
         )
+        output_price_per_million_tokens = self.get_price_per_million_tokens(
+            relevant_prices, "output"
+        )
+
+        input_price_per_token = input_price_per_million_tokens / 1_000_000
+        output_price_per_token = output_price_per_million_tokens / 1_000_000
+
+        input_cost = input_tokens * input_price_per_token
+        output_cost = output_tokens * output_price_per_token
+        cost = input_cost + output_cost
         return {
-            "input_price_source": input_price_source,
-            "input_price_per_token": input_price_per_token,
+            "input_price_per_million_tokens": input_price_per_million_tokens,
+            "output_price_per_million_tokens": output_price_per_million_tokens,
             "input_tokens": input_tokens,
-            "output_price_source": output_price_source,
             "output_tokens": output_tokens,
-            "output_price_per_token": output_price_per_token,
+            "input_cost_usd": input_cost,
+            "output_cost_usd": output_cost,
             "cost_usd": cost,
         }
 
@@ -328,6 +346,26 @@ class JobsPrompts:
             "model": model,
         }
 
+    def process_token_type(self, item: dict, token_type: str) -> tuple:
+        """
+        Helper function to process a single token type (input or output) for price estimation.
+        """
+        price = item[f"estimated_{token_type}_price_per_million_tokens"]
+        tokens = item[f"estimated_{token_type}_tokens"]
+        cost = item[f"estimated_{token_type}_cost_usd"]
+
+        return (
+            (item["inference_service"], item["model"], token_type, price),
+            {
+                "inference_service": item["inference_service"],
+                "model": item["model"],
+                "token_type": token_type,
+                "price_per_million_tokens": price,
+                "tokens": tokens,
+                "cost_usd": cost,
+            },
+        )
+
     def estimate_job_cost_from_external_prices(
         self, price_lookup: dict, iterations: int = 1
     ) -> dict:
@@ -341,9 +379,9 @@ class JobsPrompts:
         - 1 token = 4 characters.
         - For each prompt, output tokens = input tokens * 0.75, rounded up to the nearest integer.
         """
-        interviews = self.interviews
+        # Collect all prompt data
         data = []
-        for interview in interviews:
+        for interview in self.interviews:
             invigilators = [
                 FetchInvigilator(interview)(question)
                 for question in self.survey.questions
@@ -354,59 +392,62 @@ class JobsPrompts:
                     **prompt_details, price_lookup=price_lookup
                 )
                 price_estimates = {
+                    "estimated_input_price_per_million_tokens": prompt_cost[
+                        "input_price_per_million_tokens"
+                    ],
+                    "estimated_output_price_per_million_tokens": prompt_cost[
+                        "output_price_per_million_tokens"
+                    ],
                     "estimated_input_tokens": prompt_cost["input_tokens"],
                     "estimated_output_tokens": prompt_cost["output_tokens"],
+                    "estimated_input_cost_usd": prompt_cost["input_cost_usd"],
+                    "estimated_output_cost_usd": prompt_cost["output_cost_usd"],
                     "estimated_cost_usd": prompt_cost["cost_usd"],
                 }
-                data.append({**price_estimates, **prompt_details})
+                data.append(
+                    {
+                        **prompt_details,
+                        **price_estimates,
+                    }
+                )
 
-        model_groups = {}
+        # Group by service, model, token type, and price
+        detailed_groups = {}
         for item in data:
-            key = (item["inference_service"], item["model"])
-            if key not in model_groups:
-                model_groups[key] = {
-                    "inference_service": item["inference_service"],
-                    "model": item["model"],
-                    "estimated_cost_usd": 0,
-                    "estimated_input_tokens": 0,
-                    "estimated_output_tokens": 0,
-                }
-
-            # Accumulate values
-            model_groups[key]["estimated_cost_usd"] += item["estimated_cost_usd"]
-            model_groups[key]["estimated_input_tokens"] += item[
-                "estimated_input_tokens"
-            ]
-            model_groups[key]["estimated_output_tokens"] += item[
-                "estimated_output_tokens"
-            ]
-
-        # Apply iterations and convert to list
-        estimated_costs_by_model = []
-        for group_data in model_groups.values():
-            group_data["estimated_cost_usd"] *= iterations
-            group_data["estimated_input_tokens"] *= iterations
-            group_data["estimated_output_tokens"] *= iterations
-            estimated_costs_by_model.append(group_data)
+            for token_type in ["input", "output"]:
+                key, group_data = self.process_token_type(item, token_type)
+                if key not in detailed_groups:
+                    detailed_groups[key] = group_data
+                else:
+                    detailed_groups[key]["tokens"] += group_data["tokens"]
+                    detailed_groups[key]["cost_usd"] += group_data["cost_usd"]
+
+        # Apply iterations and prepare final output
+        detailed_costs = []
+        for group in detailed_groups.values():
+            group["tokens"] *= iterations
+            group["cost_usd"] *= iterations
+            detailed_costs.append(group)
 
         # Calculate totals
-        estimated_total_cost = sum(
-            model["estimated_cost_usd"] for model in estimated_costs_by_model
-        )
+        estimated_total_cost = sum(group["cost_usd"] for group in detailed_costs)
         estimated_total_input_tokens = sum(
-            model["estimated_input_tokens"] for model in estimated_costs_by_model
+            group["tokens"]
+            for group in detailed_costs
+            if group["token_type"] == "input"
         )
         estimated_total_output_tokens = sum(
-            model["estimated_output_tokens"] for model in estimated_costs_by_model
+            group["tokens"]
+            for group in detailed_costs
+            if group["token_type"] == "output"
         )
 
         output = {
             "estimated_total_cost_usd": estimated_total_cost,
             "estimated_total_input_tokens": estimated_total_input_tokens,
             "estimated_total_output_tokens": estimated_total_output_tokens,
-            "model_costs": estimated_costs_by_model,
+            "detailed_costs": detailed_costs,
         }
-
         return output
 
     def estimate_job_cost(self, iterations: int = 1) -> dict:

edsl/jobs/jobs_remote_inference_logger.py

@@ -30,6 +30,8 @@ class JobsInfo:
     error_report_url: str = None
     results_uuid: str = None
     results_url: str = None
+    completed_interviews: int = None
+    failed_interviews: int = None
 
     pretty_names = {
         "job_uuid": "Job UUID",
@@ -53,6 +55,8 @@ class JobLogger(ABC):
             "error_report_url",
             "results_uuid",
             "results_url",
+            "completed_interviews",
+            "failed_interviews",
         ],
         value: str,
     ):

edsl/jobs/jobs_runner_status.py

@@ -10,7 +10,7 @@ from typing import Any, Dict, Optional, TYPE_CHECKING
 from uuid import UUID
 
 if TYPE_CHECKING:
-    from .jobs_runner_asyncio import JobsRunnerAsyncio
+    from .jobs import Jobs
 
 
 @dataclass
@@ -65,14 +65,14 @@ class StatisticsTracker:
 class JobsRunnerStatusBase(ABC):
     def __init__(
         self,
-        jobs_runner: "JobsRunnerAsyncio",
+        jobs: "Jobs",
         n: int,
         refresh_rate: float = 1,
         endpoint_url: Optional[str] = "http://localhost:8000",
         job_uuid: Optional[UUID] = None,
         api_key: str = None,
     ):
-        self.jobs_runner = jobs_runner
+        self.jobs = jobs
         self.job_uuid = job_uuid
         self.base_url = f"{endpoint_url}"
         self.refresh_rate = refresh_rate
@@ -86,10 +86,10 @@ class JobsRunnerStatusBase(ABC):
             "unfixed_exceptions",
             "throughput",
         ]
-        self.num_total_interviews = n * len(self.jobs_runner)
+        self.num_total_interviews = n * len(self.jobs)
 
         self.distinct_models = list(
-            set(model.model for model in self.jobs_runner.jobs.models)
+            set(model.model for model in self.jobs.models)
         )
 
         self.stats_tracker = StatisticsTracker(
@@ -151,26 +151,31 @@ class JobsRunnerStatusBase(ABC):
         }
 
         model_queues = {}
-        # for model, bucket in self.jobs_runner.bucket_collection.items():
-        for model, bucket in self.jobs_runner.environment.bucket_collection.items():
-            model_name = model.model
-            model_queues[model_name] = {
-                "language_model_name": model_name,
-                "requests_bucket": {
-                    "completed": bucket.requests_bucket.num_released,
-                    "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),
-                },
-                "tokens_bucket": {
-                    "completed": bucket.tokens_bucket.num_released,
-                    "requested": bucket.tokens_bucket.num_requests,
-                    "tokens_returned": bucket.tokens_bucket.tokens_returned,
-                    "target_rate": round(bucket.tokens_bucket.target_rate, 1),
-                    "current_rate": round(bucket.tokens_bucket.get_throughput(), 1),
-                },
-            }
+        # 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():
+                model_name = model.model
+                model_queues[model_name] = {
+                    "language_model_name": model_name,
+                    "requests_bucket": {
+                        "completed": bucket.requests_bucket.num_released,
+                        "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),
+                    },
+                    "tokens_bucket": {
+                        "completed": bucket.tokens_bucket.num_released,
+                        "requested": bucket.tokens_bucket.num_requests,
+                        "tokens_returned": bucket.tokens_bucket.tokens_returned,
+                        "target_rate": round(bucket.tokens_bucket.target_rate, 1),
+                        "current_rate": round(bucket.tokens_bucket.get_throughput(), 1),
+                    },
+                }
         status_dict["language_model_queues"] = model_queues
         return status_dict
 

edsl/jobs/progress_bar_manager.py

@@ -0,0 +1,79 @@
+"""
+Progress bar management for asynchronous job execution.
+
+This module provides a context manager for handling progress bar setup and thread
+management during job execution. It coordinates the display and updating of progress
+bars, particularly for remote tracking via the Expected Parrot API.
+"""
+
+import threading
+import warnings
+
+from ..coop import Coop
+from .jobs_runner_status import JobsRunnerStatus
+
+
+class ProgressBarManager:
+    """Context manager for handling progress bar setup and thread management.
+    
+    This class manages the progress bar display and updating during job execution,
+    particularly for remote tracking via the Expected Parrot API.
+    
+    It handles:
+    1. Setting up a status tracking object
+    2. Creating and managing a background thread for progress updates
+    3. Properly cleaning up resources when execution completes
+    """
+    
+    def __init__(self, jobs, run_config, parameters):
+        self.parameters = parameters
+        self.jobs = jobs
+        
+        # Set up progress tracking
+        coop = Coop()
+        endpoint_url = coop.get_progress_bar_url()
+        
+        # Set up jobs status object
+        params = {
+            "jobs": jobs,
+            "n": parameters.n,
+            "endpoint_url": endpoint_url,
+            "job_uuid": parameters.job_uuid,
+        }
+        
+        # If the jobs_runner_status is already set, use it directly
+        if run_config.environment.jobs_runner_status is not None:
+            self.jobs_runner_status = run_config.environment.jobs_runner_status
+        else:
+            # Otherwise create a new one
+            self.jobs_runner_status = JobsRunnerStatus(**params)
+        
+        # Store on run_config for use by other components
+        run_config.environment.jobs_runner_status = self.jobs_runner_status
+        
+        self.progress_thread = None
+        self.stop_event = threading.Event()
+
+    def __enter__(self):
+        if self.parameters.progress_bar and self.jobs_runner_status.has_ep_api_key():
+            self.jobs_runner_status.setup()
+            self.progress_thread = threading.Thread(
+                target=self._run_progress_bar,
+                args=(self.stop_event, self.jobs_runner_status)
+            )
+            self.progress_thread.start()
+        elif self.parameters.progress_bar:
+            warnings.warn(
+                "You need an Expected Parrot API key to view job progress bars."
+            )
+        return self.stop_event
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop_event.set()
+        if self.progress_thread is not None:
+            self.progress_thread.join()
+
+    @staticmethod
+    def _run_progress_bar(stop_event, jobs_runner_status):
+        """Runs the progress bar in a separate thread."""
+        jobs_runner_status.update_progress(stop_event)

edsl/jobs/remote_inference.py

@@ -1,3 +1,4 @@
+import re
 from typing import Optional, Union, Literal, TYPE_CHECKING, NewType, Callable, Any
 from dataclasses import dataclass
 from ..coop import CoopServerResponseError
@@ -112,13 +113,18 @@ class JobsRemoteInferenceHandler:
         )
         logger.add_info("job_uuid", job_uuid)
 
+        remote_inference_url = self.remote_inference_url
+        if "localhost" in remote_inference_url:
+            remote_inference_url = remote_inference_url.replace("8000", "1234")
         logger.update(
-            f"Job details are available at your Coop account. [Go to Remote Inference page]({self.remote_inference_url})",
+            f"Job details are available at your Coop account. [Go to Remote Inference page]({remote_inference_url})",
             status=JobsStatus.RUNNING,
         )
         progress_bar_url = (
             f"{self.expected_parrot_url}/home/remote-job-progress/{job_uuid}"
         )
+        if "localhost" in progress_bar_url:
+            progress_bar_url = progress_bar_url.replace("8000", "1234")
         logger.add_info("progress_bar_url", progress_bar_url)
         logger.update(
             f"View job progress [here]({progress_bar_url})", status=JobsStatus.RUNNING
@@ -200,10 +206,35 @@ class JobsRemoteInferenceHandler:
             status=JobsStatus.FAILED,
         )
 
+    def _handle_partially_failed_job_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
+
     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")
 
         if latest_error_report_url:
@@ -244,6 +275,8 @@ class JobsRemoteInferenceHandler:
         job_info.logger.add_info("results_uuid", results_uuid)
         results = object_fetcher(results_uuid, expected_object_type="results")
         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":
@@ -256,6 +289,7 @@ class JobsRemoteInferenceHandler:
                 f"View partial results [here]({results_url})",
                 status=JobsStatus.PARTIALLY_FAILED,
             )
+
         results.job_uuid = job_info.job_uuid
         results.results_uuid = results_uuid
         return results

edsl/key_management/key_lookup_builder.py

@@ -2,6 +2,7 @@ from typing import Optional, TYPE_CHECKING
 import os
 from functools import lru_cache
 import textwrap
+import requests
 
 if TYPE_CHECKING:
     from ..coop import Coop
@@ -255,7 +256,11 @@ class KeyLookupBuilder:
         return dict(list(os.environ.items()))
 
     def _coop_key_value_pairs(self):
-        return dict(list(self.coop.fetch_rate_limit_config_vars().items()))
+        try:
+            return dict(list(self.coop.fetch_rate_limit_config_vars().items()))
+        except requests.ConnectionError:
+            # If connection fails, return empty dict instead of raising error
+            return {}
 
     def _config_key_value_pairs(self):
         from ..config import CONFIG