edsl 0.1.38.dev4__py3-none-any.whl → 0.1.39__py3-none-any.whl

edsl/jobs/runners/JobsRunnerStatus.py

@@ -3,21 +3,12 @@ from __future__ import annotations
 import os
 import time
 import requests
-import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-
-from typing import Any, List, DefaultDict, Optional, Dict
 from collections import defaultdict
+from typing import Any, Dict, Optional
 from uuid import UUID
 
-from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
-
-InterviewTokenUsageMapping = DefaultDict[str, InterviewTokenUsage]
-
-from edsl.jobs.interviews.InterviewStatistic import InterviewStatistic
-from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
-
 
 @dataclass
 class ModelInfo:
@@ -28,11 +19,44 @@ class ModelInfo:
     token_usage_info: dict
 
 
-@dataclass
-class ModelTokenUsageStats:
-    token_usage_type: str
-    details: List[dict]
-    cost: str
+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):
@@ -46,48 +70,39 @@ class JobsRunnerStatusBase(ABC):
         api_key: str = None,
     ):
         self.jobs_runner = jobs_runner
-
-        # The uuid of the job on Coop
         self.job_uuid = job_uuid
-
         self.base_url = f"{endpoint_url}"
-
-        self.start_time = time.time()
-        self.completed_interviews = []
         self.refresh_rate = refresh_rate
         self.statistics = [
             "elapsed_time",
             "total_interviews_requested",
             "completed_interviews",
-            #            "percent_complete",
             "average_time_per_interview",
-            #            "task_remaining",
             "estimated_time_remaining",
             "exceptions",
             "unfixed_exceptions",
             "throughput",
         ]
-        self.num_total_interviews = n * len(self.jobs_runner.interviews)
+        self.num_total_interviews = n * len(self.jobs_runner)
 
         self.distinct_models = list(
-            set(i.model.model for i in self.jobs_runner.interviews)
+            set(model.model for model in self.jobs_runner.jobs.models)
         )
 
-        self.completed_interview_by_model = defaultdict(list)
+        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.
-        """
+        """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.
-        """
+        """Converts current status into a JSON-serializable dictionary."""
         # Get all statistics
         stats = {}
         for stat_name in self.statistics:
@@ -95,42 +110,46 @@ class JobsRunnerStatusBase(ABC):
             name, value = list(stat.items())[0]
             stats[name] = value
 
-        # Calculate overall progress
-        total_interviews = len(self.jobs_runner.total_interviews)
-        completed = len(self.completed_interviews)
-
         # 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_for_model = len(self.completed_interview_by_model[model])
-            target_for_model = int(
-                self.num_total_interviews / len(self.distinct_models)
-            )
+            completed = self.stats_tracker.completed_by_model[model]
             model_progress[model] = {
-                "completed": completed_for_model,
-                "total": target_for_model,
+                "completed": completed,
+                "total": target_per_model,
                 "percent": (
-                    (completed_for_model / target_for_model * 100)
-                    if target_for_model > 0
-                    else 0
+                    (completed / target_per_model * 100) if target_per_model > 0 else 0
                 ),
             }
 
         status_dict = {
             "overall_progress": {
-                "completed": completed,
-                "total": total_interviews,
+                "completed": self.stats_tracker.completed_count,
+                "total": self.num_total_interviews,
                 "percent": (
-                    (completed / total_interviews * 100) if total_interviews > 0 else 0
+                    (
+                        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 completed >= total_interviews else "running",
+            "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.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,
@@ -152,99 +171,68 @@ class JobsRunnerStatusBase(ABC):
         status_dict["language_model_queues"] = model_queues
         return status_dict
 
-    @abstractmethod
-    def setup(self):
-        """
-        Conducts any setup that needs to happen prior to sending status updates.
+    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
+            ),
+        )
 
-        Ex. For a local job, creates a job in the Coop database.
-        """
-        pass
+    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.")}
 
-    @abstractmethod
-    def send_status_update(self):
-        """
-        Updates the current status of the job.
-        """
-        pass
+        elif stat_name == "total_interviews_requested":
+            return {"total_interviews_requested": (self.num_total_interviews, None, "")}
 
-    def add_completed_interview(self, result):
-        self.completed_interviews.append(result.interview_hash)
+        elif stat_name == "completed_interviews":
+            return {
+                "completed_interviews": (self.stats_tracker.completed_count, None, "")
+            }
 
-        relevant_model = result.model.model
-        self.completed_interview_by_model[relevant_model].append(result.interview_hash)
+        elif stat_name == "average_time_per_interview":
+            value = self.stats_tracker.get_average_time_per_interview()
+            return {"average_time_per_interview": (value, 2, "sec.")}
 
-    def _compute_statistic(self, stat_name: str):
-        completed_tasks = self.completed_interviews
-        elapsed_time = time.time() - self.start_time
-        interviews = self.jobs_runner.total_interviews
+        elif stat_name == "estimated_time_remaining":
+            value = self.stats_tracker.get_estimated_time_remaining()
+            return {"estimated_time_remaining": (value, 1, "sec.")}
 
-        stat_definitions = {
-            "elapsed_time": lambda: InterviewStatistic(
-                "elapsed_time", value=elapsed_time, digits=1, units="sec."
-            ),
-            "total_interviews_requested": lambda: InterviewStatistic(
-                "total_interviews_requested", value=len(interviews), units=""
-            ),
-            "completed_interviews": lambda: InterviewStatistic(
-                "completed_interviews", value=len(completed_tasks), units=""
-            ),
-            "percent_complete": lambda: InterviewStatistic(
-                "percent_complete",
-                value=(
-                    len(completed_tasks) / len(interviews) * 100
-                    if len(interviews) > 0
-                    else 0
-                ),
-                digits=1,
-                units="%",
-            ),
-            "average_time_per_interview": lambda: InterviewStatistic(
-                "average_time_per_interview",
-                value=elapsed_time / len(completed_tasks) if completed_tasks else 0,
-                digits=2,
-                units="sec.",
-            ),
-            "task_remaining": lambda: InterviewStatistic(
-                "task_remaining", value=len(interviews) - len(completed_tasks), units=""
-            ),
-            "estimated_time_remaining": lambda: InterviewStatistic(
-                "estimated_time_remaining",
-                value=(
-                    (len(interviews) - len(completed_tasks))
-                    * (elapsed_time / len(completed_tasks))
-                    if len(completed_tasks) > 0
-                    else 0
-                ),
-                digits=1,
-                units="sec.",
-            ),
-            "exceptions": lambda: InterviewStatistic(
-                "exceptions",
-                value=sum(len(i.exceptions) for i in interviews),
-                units="",
-            ),
-            "unfixed_exceptions": lambda: InterviewStatistic(
-                "unfixed_exceptions",
-                value=sum(i.exceptions.num_unfixed() for i in interviews),
-                units="",
-            ),
-            "throughput": lambda: InterviewStatistic(
-                "throughput",
-                value=len(completed_tasks) / elapsed_time if elapsed_time > 0 else 0,
-                digits=2,
-                units="interviews/sec.",
-            ),
-        }
-        return stat_definitions[stat_name]()
+        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
@@ -260,49 +248,35 @@ class JobsRunnerStatus(JobsRunnerStatusBase):
         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"}
-
-        if self.api_key:
-            headers["Authorization"] = f"Bearer {self.api_key}"
-        else:
-            headers["Authorization"] = f"Bearer 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:
-            # Create a new local job
             response = requests.post(
                 self.create_url,
                 headers=headers,
                 timeout=1,
             )
-        response.raise_for_status()
-        data = response.json()
-        self.job_uuid = data.get("job_uuid")
+            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.
-        """
+        """Sends current status to the web endpoint using the instance's job_uuid."""
         try:
-            # Get the status dictionary and add the job_id
             status_dict = self.get_status_dict()
-
-            # Make the UUID JSON serializable
             status_dict["job_id"] = str(self.job_uuid)
 
-            headers = {"Content-Type": "application/json"}
-
-            if self.api_key:
-                headers["Authorization"] = f"Bearer {self.api_key}"
-            else:
-                headers["Authorization"] = f"Bearer None"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.api_key or 'None'}",
+            }
 
-            # Send the update
             response = requests.patch(
                 self.update_url,
                 json=status_dict,
@@ -314,14 +288,8 @@ class JobsRunnerStatus(JobsRunnerStatusBase):
             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. Otherwise, returns False.
-        """
-
-        if self.api_key is not None:
-            return True
-        else:
-            return False
+        """Returns True if the user has an Expected Parrot API key."""
+        return self.api_key is not None
 
 
 if __name__ == "__main__":

edsl/jobs/tasks/QuestionTaskCreator.py

@@ -1,17 +1,17 @@
 import asyncio
-from typing import Callable, Union, List
+from typing import Callable, Union, List, TYPE_CHECKING
 from collections import UserList, UserDict
 
-from edsl.jobs.buckets import ModelBuckets
-from edsl.exceptions import InterviewErrorPriorTaskCanceled
+from edsl.exceptions.jobs import InterviewErrorPriorTaskCanceled
 
-from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
 from edsl.jobs.tasks.task_status_enum import TaskStatus, TaskStatusDescriptor
 from edsl.jobs.tasks.TaskStatusLog import TaskStatusLog
-from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
 from edsl.jobs.tokens.TokenUsage import TokenUsage
 from edsl.jobs.Answers import Answers
-from edsl.questions.QuestionBase import QuestionBase
+
+if TYPE_CHECKING:
+    from edsl.questions.QuestionBase import QuestionBase
+    from edsl.jobs.buckets import ModelBuckets
 
 
 class TokensUsed(UserDict):
@@ -24,7 +24,6 @@ class TokensUsed(UserDict):
 
 class QuestionTaskCreator(UserList):
     """Class to create and manage a single question and its dependencies.
-    The class is an instance of a UserList of tasks that must be completed before the focal task can be run.
 
     It is a UserList with all the tasks that must be completed before the focal task can be run.
     The focal task is the question that we are interested in answering.
@@ -35,9 +34,9 @@ class QuestionTaskCreator(UserList):
     def __init__(
         self,
         *,
-        question: QuestionBase,
+        question: "QuestionBase",
         answer_question_func: Callable,
-        model_buckets: ModelBuckets,
+        model_buckets: "ModelBuckets",
         token_estimator: Union[Callable, None] = None,
         iteration: int = 0,
     ):
@@ -51,14 +50,15 @@ class QuestionTaskCreator(UserList):
 
         """
         super().__init__([])
-        # answer_question_func is the 'interview.answer_question_and_record_task" method
         self.answer_question_func = answer_question_func
         self.question = question
         self.iteration = iteration
 
         self.model_buckets = model_buckets
+
         self.requests_bucket = self.model_buckets.requests_bucket
         self.tokens_bucket = self.model_buckets.tokens_bucket
+
         self.status_log = TaskStatusLog()
 
         def fake_token_estimator(question):
@@ -125,11 +125,13 @@ class QuestionTaskCreator(UserList):
 
         await self.tokens_bucket.get_tokens(requested_tokens)
 
-        if (estimated_wait_time := self.requests_bucket.wait_time(1)) > 0:
+        if (estimated_wait_time := self.model_buckets.requests_bucket.wait_time(1)) > 0:
             self.waiting = True  #  do we need this?
             self.task_status = TaskStatus.WAITING_FOR_REQUEST_CAPACITY
 
-        await self.requests_bucket.get_tokens(1, cheat_bucket_capacity=True)
+        await self.model_buckets.requests_bucket.get_tokens(
+            1, cheat_bucket_capacity=True
+        )
 
         self.task_status = TaskStatus.API_CALL_IN_PROGRESS
         try:
@@ -142,22 +144,22 @@ class QuestionTaskCreator(UserList):
             raise e
 
         if results.cache_used:
-            self.tokens_bucket.add_tokens(requested_tokens)
-            self.requests_bucket.add_tokens(1)
+            self.model_buckets.tokens_bucket.add_tokens(requested_tokens)
+            self.model_buckets.requests_bucket.add_tokens(1)
             self.from_cache = True
             # Turbo mode means that we don't wait for tokens or requests.
-            self.tokens_bucket.turbo_mode_on()
-            self.requests_bucket.turbo_mode_on()
+            self.model_buckets.tokens_bucket.turbo_mode_on()
+            self.model_buckets.requests_bucket.turbo_mode_on()
         else:
-            self.tokens_bucket.turbo_mode_off()
-            self.requests_bucket.turbo_mode_off()
+            self.model_buckets.tokens_bucket.turbo_mode_off()
+            self.model_buckets.requests_bucket.turbo_mode_off()
 
         return results
 
     @classmethod
     def example(cls):
         """Return an example instance of the class."""
-        from edsl import QuestionFreeText
+        from edsl.questions.QuestionFreeText import QuestionFreeText
         from edsl.jobs.buckets.ModelBuckets import ModelBuckets
 
         m = ModelBuckets.infinity_bucket()

edsl/jobs/tasks/TaskHistory.py

@@ -1,18 +1,17 @@
 from typing import List, Optional
 from io import BytesIO
-import webbrowser
-import os
 import base64
-from importlib import resources
 from edsl.jobs.tasks.task_status_enum import TaskStatus
+from edsl.Base import RepresentationMixin
 
 
-class TaskHistory:
+class TaskHistory(RepresentationMixin):
     def __init__(
         self,
-        interviews: List["Interview"],
+        interviews: List["Interview"] = None,
         include_traceback: bool = False,
         max_interviews: int = 10,
+        interviews_with_exceptions_only: bool = False,
     ):
         """
         The structure of a TaskHistory exception
@@ -22,13 +21,33 @@ class TaskHistory:
         >>> _ = TaskHistory.example()
         ...
         """
+        self.interviews_with_exceptions_only = interviews_with_exceptions_only
+        self._interviews = {}
+        self.total_interviews = []
+        if interviews is not None:
+            for interview in interviews:
+                self.add_interview(interview)
 
-        self.total_interviews = interviews
+        self.include_traceback = include_traceback
+        self._interviews = {
+            index: interview for index, interview in enumerate(self.total_interviews)
+        }
+        self.max_interviews = max_interviews
+
+        # self.total_interviews = interviews
         self.include_traceback = include_traceback
 
-        self._interviews = {index: i for index, i in enumerate(self.total_interviews)}
+        # self._interviews = {index: i for index, i in enumerate(self.total_interviews)}
         self.max_interviews = max_interviews
 
+    def add_interview(self, interview: "Interview"):
+        """Add a single interview to the history"""
+        if self.interviews_with_exceptions_only and interview.exceptions == {}:
+            return
+
+        self.total_interviews.append(interview)
+        self._interviews[len(self._interviews)] = interview
+
     @classmethod
     def example(cls):
         """ """
@@ -121,14 +140,6 @@ class TaskHistory:
         """Return True if there are any exceptions."""
         return len(self.unfixed_exceptions) > 0
 
-    def _repr_html_(self):
-        """Return an HTML representation of the TaskHistory."""
-        d = self.to_dict(add_edsl_version=False)
-        data = [[k, v] for k, v in d.items()]
-        from tabulate import tabulate
-
-        return tabulate(data, headers=["keys", "values"], tablefmt="html")
-
     def show_exceptions(self, tracebacks=False):
         """Print the exceptions."""
         for index in self.indices:
@@ -240,11 +251,15 @@ class TaskHistory:
             plt.show()
 
     def css(self):
+        from importlib import resources
+
         env = resources.files("edsl").joinpath("templates/error_reporting")
         css = env.joinpath("report.css").read_text()
         return css
 
     def javascript(self):
+        from importlib import resources
+
         env = resources.files("edsl").joinpath("templates/error_reporting")
         js = env.joinpath("report.js").read_text()
         return js
@@ -281,7 +296,7 @@ class TaskHistory:
         exceptions_by_question_name = {}
         for interview in self.total_interviews:
             for question_name, exceptions in interview.exceptions.items():
-                question_type = interview.survey.get_question(
+                question_type = interview.survey._get_question_by_name(
                     question_name
                 ).question_type
                 if (question_name, question_type) not in exceptions_by_question_name:
@@ -330,8 +345,11 @@ class TaskHistory:
         }
         return sorted_exceptions_by_model
 
-    def generate_html_report(self, css: Optional[str]):
-        performance_plot_html = self.plot(num_periods=100, get_embedded_html=True)
+    def generate_html_report(self, css: Optional[str], include_plot=False):
+        if include_plot:
+            performance_plot_html = self.plot(num_periods=100, get_embedded_html=True)
+        else:
+            performance_plot_html = ""
 
         if css is None:
             css = self.css()
@@ -409,6 +427,8 @@ class TaskHistory:
             print(f"Exception report saved to {filename}")
 
         if open_in_browser:
+            import webbrowser
+
             webbrowser.open(f"file://{os.path.abspath(filename)}")
 
         if return_link:

edsl/jobs/tasks/task_status_enum.py

@@ -3,8 +3,6 @@ from collections import UserDict
 import enum
 import time
 
-# from edsl.jobs.tasks.TaskStatusLogEntry import TaskStatusLogEntry
-
 
 class TaskStatus(enum.Enum):
     "These are the possible states a task can be in."