edsl 0.1.15__py3-none-any.whl → 0.1.40__py3-none-any.whl

edsl/jobs/results_exceptions_handler.py

@@ -0,0 +1,98 @@
+from typing import Optional, TYPE_CHECKING, Protocol
+import sys
+from edsl.scenarios.FileStore import HTMLFileStore
+from edsl.config import CONFIG
+from edsl.coop.coop import Coop
+
+
+class ResultsProtocol(Protocol):
+    """Protocol defining the required interface for Results objects."""
+
+    @property
+    def has_unfixed_exceptions(self) -> bool: ...
+
+    @property
+    def task_history(self) -> "TaskHistoryProtocol": ...
+
+
+class TaskHistoryProtocol(Protocol):
+    """Protocol defining the required interface for TaskHistory objects."""
+
+    @property
+    def indices(self) -> list: ...
+
+    def html(self, cta: str, open_in_browser: bool, return_link: bool) -> str: ...
+
+
+class RunParametersProtocol(Protocol):
+    """Protocol defining the required interface for RunParameters objects."""
+
+    @property
+    def print_exceptions(self) -> bool: ...
+
+
+class ResultsExceptionsHandler:
+    """Handles exception reporting and display functionality."""
+
+    def __init__(
+        self, results: ResultsProtocol, parameters: RunParametersProtocol
+    ) -> None:
+        self.results = results
+        self.parameters = parameters
+
+        self.open_in_browser = self._get_browser_setting()
+        self.remote_logging = self._get_remote_logging_setting()
+
+    def _get_browser_setting(self) -> bool:
+        """Determine if exceptions should be opened in browser based on config."""
+        setting = CONFIG.get("EDSL_OPEN_EXCEPTION_REPORT_URL")
+        if setting == "True":
+            return True
+        elif setting == "False":
+            return False
+        else:
+            raise Exception(
+                "EDSL_OPEN_EXCEPTION_REPORT_URL must be either True or False"
+            )
+
+    def _get_remote_logging_setting(self) -> bool:
+        """Get remote logging setting from coop."""
+        try:
+            coop = Coop()
+            return coop.edsl_settings["remote_logging"]
+        except Exception as e:
+            # print(e)
+            return False
+
+    def _generate_error_message(self, indices) -> str:
+        """Generate appropriate error message based on number of exceptions."""
+        msg = f"Exceptions were raised in {len(indices)} interviews.\n"
+        if len(indices) > 5:
+            msg += f"Exceptions were raised in the following interviews: {indices}.\n"
+        return msg
+
+    def handle_exceptions(self) -> None:
+        """Handle exceptions by printing messages and generating reports as needed."""
+        if not (
+            self.results.has_unfixed_exceptions and self.parameters.print_exceptions
+        ):
+            return
+
+        # Print error message
+        error_msg = self._generate_error_message(self.results.task_history.indices)
+        print(error_msg, file=sys.stderr)
+
+        # Generate HTML report
+        filepath = self.results.task_history.html(
+            cta="Open report to see details.",
+            open_in_browser=self.open_in_browser,
+            return_link=True,
+        )
+
+        # Handle remote logging if enabled
+        if self.remote_logging:
+            filestore = HTMLFileStore(filepath)
+            coop_details = filestore.push(description="Error report")
+            print(coop_details)
+
+        print("Also see: https://docs.expectedparrot.com/en/latest/exceptions.html")

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -1,122 +1,163 @@
+from __future__ import annotations
 import time
 import asyncio
-from typing import Coroutine, List, AsyncGenerator
+import threading
+import warnings
+from typing import TYPE_CHECKING
 
-from rich.live import Live
-from rich.console import Console
-
-from edsl.results import Results, Result
-from edsl.jobs.JobsRunner import JobsRunner
-from edsl.jobs.Interview import Interview
+from edsl.results.Results import Results
+from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
+from edsl.jobs.tasks.TaskHistory import TaskHistory
 from edsl.utilities.decorators import jupyter_nb_handler
+from edsl.jobs.async_interview_runner import AsyncInterviewRunner
+from edsl.jobs.data_structures import RunEnvironment, RunParameters, RunConfig
 
-from edsl.jobs.JobsRunnerStatusMixin import JobsRunnerStatusMixin
-
-
-class JobsRunnerAsyncio(JobsRunner, JobsRunnerStatusMixin):
-    runner_name = "asyncio"
-
-    async def run_async(
-        self, n=1, verbose=False, sleep=0, debug=False, progress_bar=False
-    ) -> AsyncGenerator[Result, None]:
-        """Creates the tasks, runs them asynchronously, and returns the results as a Results object.
-        Completed tasks are yielded as they are completed.
-        """
-        tasks = self._create_all_interview_tasks(self.interviews, debug)
-        for task in asyncio.as_completed(tasks):
-            result = await task
-            yield result
-
-    def _create_all_interview_tasks(self, interviews, debug) -> List[asyncio.Task]:
-        """Creates an awaitable task for each interview."""
-        tasks = []
-        for i, interview in enumerate(interviews):
-            interviewing_task = self._interview_task(interview, i, debug)
-            tasks.append(asyncio.create_task(interviewing_task))
-        return tasks
-
-    async def _interview_task(
-        self, interview: Interview, i: int, debug: bool
-    ) -> Result:
-        """Conducts an interview and returns the result."""
-        # the model buckets are used to track usage rates
-        model_buckets = self.bucket_collection[interview.model]
-
-        # get the results of the interview
-        answer, valid_results = await interview.async_conduct_interview(
-            debug=debug, model_buckets=model_buckets
-        )
-        # breakpoint()
-
-        # we should have a valid result for each question
-        answer_key_names = {k for k in set(answer.keys()) if not k.endswith("_comment")}
-        assert len(valid_results) == len(answer_key_names)
-
-        question_name_to_prompts = dict({})
-        for result in valid_results:
-            question_name = result["question_name"]
-            question_name_to_prompts[question_name] = {
-                "user_prompt": result["prompts"]["user_prompt"],
-                "system_prompt": result["prompts"]["system_prompt"],
-            }
-
-        prompt_dictionary = {}
-        for answer_key_name in answer_key_names:
-            prompt_dictionary[
-                answer_key_name + "_user_prompt"
-            ] = question_name_to_prompts[answer_key_name]["user_prompt"]
-            prompt_dictionary[
-                answer_key_name + "_system_prompt"
-            ] = question_name_to_prompts[answer_key_name]["system_prompt"]
-
-        raw_model_results_dictionary = {}
-        for result in valid_results:
-            question_name = result["question_name"]
-            raw_model_results_dictionary[
-                question_name + "_raw_model_response"
-            ] = result["raw_model_response"]
-
-        result = Result(
-            agent=interview.agent,
-            scenario=interview.scenario,
-            model=interview.model,
-            iteration=i,
-            answer=answer,
-            prompt=prompt_dictionary,
-            raw_model_response=raw_model_results_dictionary,
-        )
-        return result
+if TYPE_CHECKING:
+    from edsl.jobs.Jobs import Jobs
 
-    @jupyter_nb_handler
-    async def run(
-        self, n=1, verbose=True, sleep=0, debug=False, progress_bar=False
-    ) -> Coroutine:
-        """Runs a collection of interviews, handling both async and sync contexts."""
-        verbose = True
-        console = Console()
+
+class JobsRunnerAsyncio:
+    """A class for running a collection of interviews asynchronously.
+
+    It gets instaniated from a Jobs object.
+    The Jobs object is a collection of interviews that are to be run.
+    """
+
+    def __init__(self, jobs: "Jobs", environment: RunEnvironment):
+        self.jobs = jobs
+        self.environment = environment
+
+    def __len__(self):
+        return len(self.jobs)
+
+    async def run_async(self, parameters: RunParameters) -> Results:
+        """Used for some other modules that have a non-standard way of running interviews."""
+
+        self.environment.jobs_runner_status = JobsRunnerStatus(self, n=parameters.n)
         data = []
-        start_time = time.monotonic()
-
-        live = None
-        if progress_bar:
-            live = Live(
-                self._generate_status_table(data, 0),
-                console=console,
-                refresh_per_second=10,
-            )
-            live.__enter__()  # Manually enter the Live context
-
-        async for result in self.run_async(n, verbose, sleep, debug, progress_bar):
-            end_time = time.monotonic()
-            elapsed_time = end_time - start_time
+        task_history = TaskHistory(include_traceback=False)
+
+        run_config = RunConfig(parameters=parameters, environment=self.environment)
+        result_generator = AsyncInterviewRunner(self.jobs, run_config)
+
+        async for result, interview in result_generator.run():
             data.append(result)
+            task_history.add_interview(interview)
 
-            if progress_bar:
-                live.update(self._generate_status_table(data, elapsed_time))
+        results = Results(survey=self.jobs.survey, task_history=task_history, data=data)
 
-        if progress_bar:
-            live.update(self._generate_status_table(data, elapsed_time))
-            await asyncio.sleep(0.5)  # short delay to show the final status
-            live.__exit__(None, None, None)  # Manually exit the Live context
+        relevant_cache = results.relevant_cache(self.environment.cache)
 
+        return Results(
+            survey=self.jobs.survey,
+            task_history=task_history,
+            data=data,
+            cache=relevant_cache,
+        )
+
+    def simple_run(self):
+        data = asyncio.run(self.run_async())
         return Results(survey=self.jobs.survey, data=data)
+
+    @jupyter_nb_handler
+    async def run(self, parameters: RunParameters) -> Results:
+        """Runs a collection of interviews, handling both async and sync contexts."""
+
+        run_config = RunConfig(parameters=parameters, environment=self.environment)
+
+        self.start_time = time.monotonic()
+        self.completed = False
+
+        from edsl.coop import Coop
+
+        coop = Coop()
+        endpoint_url = coop.get_progress_bar_url()
+
+        def set_up_jobs_runner_status(jobs_runner_status):
+            if jobs_runner_status is not None:
+                return jobs_runner_status(
+                    self,
+                    n=parameters.n,
+                    endpoint_url=endpoint_url,
+                    job_uuid=parameters.job_uuid,
+                )
+            else:
+                return JobsRunnerStatus(
+                    self,
+                    n=parameters.n,
+                    endpoint_url=endpoint_url,
+                    job_uuid=parameters.job_uuid,
+                )
+
+        run_config.environment.jobs_runner_status = set_up_jobs_runner_status(
+            self.environment.jobs_runner_status
+        )
+
+        async def get_results(results) -> None:
+            """Conducted the interviews and append to the results list."""
+            result_generator = AsyncInterviewRunner(self.jobs, run_config)
+            async for result, interview in result_generator.run():
+                results.append(result)
+                results.task_history.add_interview(interview)
+
+            self.completed = True
+
+        def run_progress_bar(stop_event, jobs_runner_status) -> None:
+            """Runs the progress bar in a separate thread."""
+            jobs_runner_status.update_progress(stop_event)
+
+        def set_up_progress_bar(progress_bar: bool, jobs_runner_status):
+            progress_thread = None
+            if progress_bar and jobs_runner_status.has_ep_api_key():
+                jobs_runner_status.setup()
+                progress_thread = threading.Thread(
+                    target=run_progress_bar, args=(stop_event, jobs_runner_status)
+                )
+                progress_thread.start()
+            elif progress_bar:
+                warnings.warn(
+                    "You need an Expected Parrot API key to view job progress bars."
+                )
+            return progress_thread
+
+        results = Results(
+            survey=self.jobs.survey,
+            data=[],
+            task_history=TaskHistory(),
+            #           cache=self.environment.cache.new_entries_cache(),
+        )
+
+        stop_event = threading.Event()
+        progress_thread = set_up_progress_bar(
+            parameters.progress_bar, run_config.environment.jobs_runner_status
+        )
+
+        exception_to_raise = None
+        try:
+            await get_results(results)
+        except KeyboardInterrupt:
+            print("Keyboard interrupt received. Stopping gracefully...")
+            stop_event.set()
+        except Exception as e:
+            if parameters.stop_on_exception:
+                exception_to_raise = e
+            stop_event.set()
+        finally:
+            stop_event.set()
+            if progress_thread is not None:
+                progress_thread.join()
+
+            if exception_to_raise:
+                raise exception_to_raise
+
+            relevant_cache = results.relevant_cache(self.environment.cache)
+            results.cache = relevant_cache
+            # breakpoint()
+            results.bucket_collection = self.environment.bucket_collection
+
+            from edsl.jobs.results_exceptions_handler import ResultsExceptionsHandler
+
+            results_exceptions_handler = ResultsExceptionsHandler(results, parameters)
+
+            results_exceptions_handler.handle_exceptions()
+            return results

edsl/jobs/runners/JobsRunnerStatus.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+import os
+import time
+import requests
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from collections import defaultdict
+from typing import Any, Dict, Optional
+from uuid import UUID
+
+
+@dataclass
+class ModelInfo:
+    model_name: str
+    TPM_limit_k: float
+    RPM_limit_k: float
+    num_tasks_waiting: int
+    token_usage_info: dict
+
+
+class StatisticsTracker:
+    def __init__(self, total_interviews: int, distinct_models: list[str]):
+        self.start_time = time.time()
+        self.total_interviews = total_interviews
+        self.completed_count = 0
+        self.completed_by_model = defaultdict(int)
+        self.distinct_models = distinct_models
+        self.total_exceptions = 0
+        self.unfixed_exceptions = 0
+
+    def add_completed_interview(
+        self, model: str, 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
+
+    def get_elapsed_time(self) -> float:
+        return time.time() - self.start_time
+
+    def get_average_time_per_interview(self) -> float:
+        return (
+            self.get_elapsed_time() / self.completed_count
+            if self.completed_count > 0
+            else 0
+        )
+
+    def get_throughput(self) -> float:
+        elapsed = self.get_elapsed_time()
+        return self.completed_count / elapsed if elapsed > 0 else 0
+
+    def get_estimated_time_remaining(self) -> float:
+        if self.completed_count == 0:
+            return 0
+        avg_time = self.get_average_time_per_interview()
+        remaining = self.total_interviews - self.completed_count
+        return avg_time * remaining
+
+
+class JobsRunnerStatusBase(ABC):
+    def __init__(
+        self,
+        jobs_runner: "JobsRunnerAsyncio",
+        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.job_uuid = job_uuid
+        self.base_url = f"{endpoint_url}"
+        self.refresh_rate = refresh_rate
+        self.statistics = [
+            "elapsed_time",
+            "total_interviews_requested",
+            "completed_interviews",
+            "average_time_per_interview",
+            "estimated_time_remaining",
+            "exceptions",
+            "unfixed_exceptions",
+            "throughput",
+        ]
+        self.num_total_interviews = n * len(self.jobs_runner)
+
+        self.distinct_models = list(
+            set(model.model for model in self.jobs_runner.jobs.models)
+        )
+
+        self.stats_tracker = StatisticsTracker(
+            total_interviews=self.num_total_interviews,
+            distinct_models=self.distinct_models,
+        )
+
+        self.api_key = api_key or os.getenv("EXPECTED_PARROT_API_KEY")
+
+    @abstractmethod
+    def has_ep_api_key(self):
+        """Checks if the user has an Expected Parrot API key."""
+        pass
+
+    def get_status_dict(self) -> Dict[str, Any]:
+        """Converts current status into a JSON-serializable dictionary."""
+        # Get all statistics
+        stats = {}
+        for stat_name in self.statistics:
+            stat = self._compute_statistic(stat_name)
+            name, value = list(stat.items())[0]
+            stats[name] = value
+
+        # Get model-specific progress
+        model_progress = {}
+        target_per_model = int(self.num_total_interviews / len(self.distinct_models))
+
+        for model in self.distinct_models:
+            completed = self.stats_tracker.completed_by_model[model]
+            model_progress[model] = {
+                "completed": completed,
+                "total": target_per_model,
+                "percent": (
+                    (completed / target_per_model * 100) if target_per_model > 0 else 0
+                ),
+            }
+
+        status_dict = {
+            "overall_progress": {
+                "completed": self.stats_tracker.completed_count,
+                "total": self.num_total_interviews,
+                "percent": (
+                    (
+                        self.stats_tracker.completed_count
+                        / self.num_total_interviews
+                        * 100
+                    )
+                    if self.num_total_interviews > 0
+                    else 0
+                ),
+            },
+            "language_model_progress": model_progress,
+            "statistics": stats,
+            "status": (
+                "completed"
+                if self.stats_tracker.completed_count >= self.num_total_interviews
+                else "running"
+            ),
+        }
+
+        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),
+                },
+            }
+        status_dict["language_model_queues"] = model_queues
+        return status_dict
+
+    def add_completed_interview(self, result):
+        """Records a completed interview without storing the full interview data."""
+        self.stats_tracker.add_completed_interview(
+            model=result.model.model,
+            num_exceptions=(
+                len(result.exceptions) if hasattr(result, "exceptions") else 0
+            ),
+            num_unfixed=(
+                result.exceptions.num_unfixed() if hasattr(result, "exceptions") else 0
+            ),
+        )
+
+    def _compute_statistic(self, stat_name: str):
+        """Computes individual statistics based on the stats tracker."""
+        if stat_name == "elapsed_time":
+            value = self.stats_tracker.get_elapsed_time()
+            return {"elapsed_time": (value, 1, "sec.")}
+
+        elif stat_name == "total_interviews_requested":
+            return {"total_interviews_requested": (self.num_total_interviews, None, "")}
+
+        elif stat_name == "completed_interviews":
+            return {
+                "completed_interviews": (self.stats_tracker.completed_count, None, "")
+            }
+
+        elif stat_name == "average_time_per_interview":
+            value = self.stats_tracker.get_average_time_per_interview()
+            return {"average_time_per_interview": (value, 2, "sec.")}
+
+        elif stat_name == "estimated_time_remaining":
+            value = self.stats_tracker.get_estimated_time_remaining()
+            return {"estimated_time_remaining": (value, 1, "sec.")}
+
+        elif stat_name == "exceptions":
+            return {"exceptions": (self.stats_tracker.total_exceptions, None, "")}
+
+        elif stat_name == "unfixed_exceptions":
+            return {
+                "unfixed_exceptions": (self.stats_tracker.unfixed_exceptions, None, "")
+            }
+
+        elif stat_name == "throughput":
+            value = self.stats_tracker.get_throughput()
+            return {"throughput": (value, 2, "interviews/sec.")}
+
+    def update_progress(self, stop_event):
+        while not stop_event.is_set():
+            self.send_status_update()
+            time.sleep(self.refresh_rate)
+        self.send_status_update()
+
+    @abstractmethod
+    def setup(self):
+        """Conducts any setup needed prior to sending status updates."""
+        pass
+
+    @abstractmethod
+    def send_status_update(self):
+        """Updates the current status of the job."""
+        pass
+
+
+class JobsRunnerStatus(JobsRunnerStatusBase):
+    @property
+    def create_url(self) -> str:
+        return f"{self.base_url}/api/v0/local-job"
+
+    @property
+    def viewing_url(self) -> str:
+        return f"{self.base_url}/home/local-job-progress/{str(self.job_uuid)}"
+
+    @property
+    def update_url(self) -> str:
+        return f"{self.base_url}/api/v0/local-job/{str(self.job_uuid)}"
+
+    def setup(self) -> None:
+        """Creates a local job on Coop if one does not already exist."""
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key or 'None'}",
+        }
+
+        if self.job_uuid is None:
+            response = requests.post(
+                self.create_url,
+                headers=headers,
+                timeout=1,
+            )
+            response.raise_for_status()
+            data = response.json()
+            self.job_uuid = data.get("job_uuid")
+
+        print(f"Running with progress bar. View progress at {self.viewing_url}")
+
+    def send_status_update(self) -> None:
+        """Sends current status to the web endpoint using the instance's job_uuid."""
+        try:
+            status_dict = self.get_status_dict()
+            status_dict["job_id"] = str(self.job_uuid)
+
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.api_key or 'None'}",
+            }
+
+            response = requests.patch(
+                self.update_url,
+                json=status_dict,
+                headers=headers,
+                timeout=1,
+            )
+            response.raise_for_status()
+        except requests.exceptions.RequestException as e:
+            print(f"Failed to send status update for job {self.job_uuid}: {e}")
+
+    def has_ep_api_key(self) -> bool:
+        """Returns True if the user has an Expected Parrot API key."""
+        return self.api_key is not None
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)