edsl 0.1.53__py3-none-any.whl → 0.1.55__py3-none-any.whl

edsl/jobs/jobs_pricing_estimation.py

@@ -1,7 +1,7 @@
 import logging
 import math
 
-from typing import List, TYPE_CHECKING
+from typing import List, TYPE_CHECKING, Union, Literal
 
 if TYPE_CHECKING:
     from .jobs import Jobs
@@ -26,53 +26,104 @@ class PromptCostEstimator:
     OUTPUT_TOKENS_PER_INPUT_TOKEN = 0.75
     PIPING_MULTIPLIER = 2
 
-    def __init__(self,      
+    def __init__(
+        self,
         system_prompt: str,
         user_prompt: str,
         price_lookup: dict,
         inference_service: str,
-        model: str):
+        model: str,
+    ):
         self.system_prompt = system_prompt
         self.user_prompt = user_prompt
         self.price_lookup = price_lookup
         self.inference_service = inference_service
         self.model = model
 
-    @staticmethod    
+    @staticmethod
     def get_piping_multiplier(prompt: str):
         """Returns 2 if a prompt includes Jinja braces, and 1 otherwise."""
 
         if "{{" in prompt and "}}" in prompt:
             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 input_price_per_token(self):
+
+    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).
+
+        Args:
+            price_type: Either "input" or "output"
+
+        Returns:
+            float | None: The highest price per token for the service, or None if not found
+        """
+        prices_for_service = [
+            prices[price_type]["service_stated_token_price"]
+            / prices[price_type]["service_stated_token_qty"]
+            for (service, _), prices in self.price_lookup.items()
+            if service == self.inference_service and price_type in prices
+        ]
+        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"]
+            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
+
             warnings.warn(
-                "Price data could not be retrieved. Using default estimates for input and output token prices. Input: $1.00 / 1M tokens; Output: $1.00 / 1M tokens"
+                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
+            return self.DEFAULT_INPUT_PRICE_PER_TOKEN, "default"
 
-    def output_price_per_token(self):
+    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"]
+            return (
+                self.relevant_prices["output"]["service_stated_token_price"]
+                / self.relevant_prices["output"]["service_stated_token_qty"]
+            ), "price_lookup"
         except KeyError:
-            return self.DEFAULT_OUTPUT_PRICE_PER_TOKEN
-        
+            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
+
+            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"
+
     def __call__(self):
         user_prompt_chars = len(str(self.user_prompt)) * self.get_piping_multiplier(
             str(self.user_prompt)
@@ -84,20 +135,37 @@ 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()
+
         cost = (
-            input_tokens * self.input_price_per_token()
-            + output_tokens * self.output_price_per_token()
+            input_tokens * input_price_per_token
+            + output_tokens * output_price_per_token
         )
         return {
+            "input_price_source": input_price_source,
+            "input_price_per_token": input_price_per_token,
             "input_tokens": input_tokens,
+            "output_price_source": output_price_source,
             "output_tokens": output_tokens,
+            "output_price_per_token": output_price_per_token,
             "cost_usd": cost,
         }
 
 
 class JobsPrompts:
 
-    relevant_keys = ["user_prompt", "system_prompt", "interview_index", "question_name", "scenario_index", "agent_index", "model", "estimated_cost", "cache_keys"]
+    relevant_keys = [
+        "user_prompt",
+        "system_prompt",
+        "interview_index",
+        "question_name",
+        "scenario_index",
+        "agent_index",
+        "model",
+        "estimated_cost",
+        "cache_keys",
+    ]
 
     """This generates the prompts for a job for price estimation purposes. 
     
@@ -105,7 +173,6 @@ class JobsPrompts:
     So assumptions are made about expansion of Jinja braces, etc.
     """
 
-
     @classmethod
     def from_jobs(cls, jobs: "Jobs"):
         """Construct a JobsPrompts object from a Jobs object."""
@@ -114,13 +181,16 @@ class JobsPrompts:
         scenarios = jobs.scenarios
         survey = jobs.survey
         return cls(
-            interviews=interviews,
-            agents=agents,
-            scenarios=scenarios,
-            survey=survey
+            interviews=interviews, agents=agents, scenarios=scenarios, survey=survey
         )
-    
-    def __init__(self, interviews: List['Interview'], agents:'AgentList', scenarios: 'ScenarioList', survey: 'Survey'):
+
+    def __init__(
+        self,
+        interviews: List["Interview"],
+        agents: "AgentList",
+        scenarios: "ScenarioList",
+        survey: "Survey",
+    ):
         """Initialize with extracted components rather than a Jobs object."""
         self.interviews = interviews
         self.agents = agents
@@ -143,17 +213,19 @@ class JobsPrompts:
             self._price_lookup = c.fetch_prices()
         return self._price_lookup
 
-    def _process_one_invigilator(self, invigilator: 'Invigilator', interview_index: int, iterations: int = 1) -> dict:
+    def _process_one_invigilator(
+        self, invigilator: "Invigilator", interview_index: int, iterations: int = 1
+    ) -> dict:
         """Process a single invigilator and return a dictionary with all needed data fields."""
         prompts = invigilator.get_prompts()
         user_prompt = prompts["user_prompt"]
         system_prompt = prompts["system_prompt"]
-        
+
         agent_index = self._agent_lookup[invigilator.agent]
         scenario_index = self._scenario_lookup[invigilator.scenario]
         model = invigilator.model.model
         question_name = invigilator.question.question_name
-        
+
         # Calculate prompt cost
         prompt_cost = self.estimate_prompt_cost(
             system_prompt=system_prompt,
@@ -163,7 +235,7 @@ class JobsPrompts:
             model=model,
         )
         cost = prompt_cost["cost_usd"]
-        
+
         # Generate cache keys for each iteration
         cache_keys = []
         for iteration in range(iterations):
@@ -175,7 +247,7 @@ class JobsPrompts:
                 iteration=iteration,
             )
             cache_keys.append(cache_key)
-        
+
         d = {
             "user_prompt": user_prompt,
             "system_prompt": system_prompt,
@@ -200,7 +272,7 @@ class JobsPrompts:
         dataset_of_prompts = {k: [] for k in self.relevant_keys}
 
         interviews = self.interviews
-        
+
         # Process each interview and invigilator
         for interview_index, interview in enumerate(interviews):
             invigilators = [
@@ -210,11 +282,13 @@ class JobsPrompts:
 
             for invigilator in invigilators:
                 # Process the invigilator and get all data as a dictionary
-                data = self._process_one_invigilator(invigilator, interview_index, iterations)
+                data = self._process_one_invigilator(
+                    invigilator, interview_index, iterations
+                )
                 for k in self.relevant_keys:
                     dataset_of_prompts[k].append(data[k])
-                
-        return Dataset([{k:dataset_of_prompts[k]} for k in self.relevant_keys])
+
+        return Dataset([{k: dataset_of_prompts[k]} for k in self.relevant_keys])
 
     @staticmethod
     def estimate_prompt_cost(
@@ -230,13 +304,13 @@ class JobsPrompts:
             user_prompt=user_prompt,
             price_lookup=price_lookup,
             inference_service=inference_service,
-            model=model
+            model=model,
         )()
-    
+
     @staticmethod
     def _extract_prompt_details(invigilator: FetchInvigilator) -> dict:
         """Extracts the prompt details from the invigilator.
-        
+
         >>> from edsl.invigilators import InvigilatorAI
         >>> invigilator = InvigilatorAI.example()
         >>> JobsPrompts._extract_prompt_details(invigilator)
@@ -276,11 +350,13 @@ class JobsPrompts:
             ]
             for invigilator in invigilators:
                 prompt_details = self._extract_prompt_details(invigilator)
-                prompt_cost = self.estimate_prompt_cost(**prompt_details, price_lookup=price_lookup)
+                prompt_cost = self.estimate_prompt_cost(
+                    **prompt_details, price_lookup=price_lookup
+                )
                 price_estimates = {
-                    'estimated_input_tokens': prompt_cost['input_tokens'],
-                    'estimated_output_tokens': prompt_cost['output_tokens'],
-                    'estimated_cost_usd': prompt_cost['cost_usd']
+                    "estimated_input_tokens": prompt_cost["input_tokens"],
+                    "estimated_output_tokens": prompt_cost["output_tokens"],
+                    "estimated_cost_usd": prompt_cost["cost_usd"],
                 }
                 data.append({**price_estimates, **prompt_details})
 
@@ -293,14 +369,18 @@ class JobsPrompts:
                     "model": item["model"],
                     "estimated_cost_usd": 0,
                     "estimated_input_tokens": 0,
-                    "estimated_output_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"]
-        
+            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():
@@ -345,4 +425,5 @@ class JobsPrompts:
 
 if __name__ == "__main__":
     import doctest
+
     doctest.testmod(optionflags=doctest.ELLIPSIS)

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