edsl 0.1.38__py3-none-any.whl → 0.1.38.dev1__py3-none-any.whl

edsl/jobs/runners/JobsRunnerStatus.py

@@ -1,21 +1,33 @@
 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 uuid import UUID
-
+from dataclasses import dataclass, asdict
+
+from typing import List, DefaultDict, Optional, Type, Literal
+from collections import UserDict, defaultdict
+
+from rich.text import Text
+from rich.box import SIMPLE
+from rich.table import Table
+from rich.live import Live
+from rich.panel import Panel
+from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn
+from rich.layout import Layout
+from rich.console import Group
+from rich import box
+
+from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
 from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
+from edsl.jobs.tokens.TokenUsage import TokenUsage
+from edsl.enums import get_token_pricing
+from edsl.jobs.tasks.task_status_enum import TaskStatus
 
 InterviewTokenUsageMapping = DefaultDict[str, InterviewTokenUsage]
 
 from edsl.jobs.interviews.InterviewStatistic import InterviewStatistic
+from edsl.jobs.interviews.InterviewStatisticsCollection import (
+    InterviewStatisticsCollection,
+)
 from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
 
 
@@ -35,23 +47,16 @@ class ModelTokenUsageStats:
     cost: str
 
 
-class JobsRunnerStatusBase(ABC):
+class Stats:
+    def elapsed_time(self):
+        InterviewStatistic("elapsed_time", value=elapsed_time, digits=1, units="sec.")
+
+
+class JobsRunnerStatus:
     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: "JobsRunnerAsyncio", n: int, refresh_rate: float = 0.25
     ):
         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
@@ -75,99 +80,6 @@ class JobsRunnerStatusBase(ABC):
 
         self.completed_interview_by_model = defaultdict(list)
 
-        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
-
-        # Calculate overall progress
-        total_interviews = len(self.jobs_runner.total_interviews)
-        completed = len(self.completed_interviews)
-
-        # Get model-specific progress
-        model_progress = {}
-        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)
-            )
-            model_progress[model] = {
-                "completed": completed_for_model,
-                "total": target_for_model,
-                "percent": (
-                    (completed_for_model / target_for_model * 100)
-                    if target_for_model > 0
-                    else 0
-                ),
-            }
-
-        status_dict = {
-            "overall_progress": {
-                "completed": completed,
-                "total": total_interviews,
-                "percent": (
-                    (completed / total_interviews * 100) if total_interviews > 0 else 0
-                ),
-            },
-            "language_model_progress": model_progress,
-            "statistics": stats,
-            "status": "completed" if completed >= total_interviews else "running",
-        }
-
-        model_queues = {}
-        for model, bucket in self.jobs_runner.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
-
-    @abstractmethod
-    def setup(self):
-        """
-        Conducts any setup that needs to happen prior to sending status updates.
-
-        Ex. For a local job, creates a job in the Coop database.
-        """
-        pass
-
-    @abstractmethod
-    def send_status_update(self):
-        """
-        Updates the current status of the job.
-        """
-        pass
-
     def add_completed_interview(self, result):
         self.completed_interviews.append(result.interview_hash)
 
@@ -238,90 +150,180 @@ class JobsRunnerStatusBase(ABC):
         }
         return stat_definitions[stat_name]()
 
-    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()
-
-
-class JobsRunnerStatus(JobsRunnerStatusBase):
-    @property
-    def create_url(self) -> str:
-        return f"{self.base_url}/api/v0/local-job"
+    def create_progress_bar(self):
+        return Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+            TextColumn("{task.completed}/{task.total}"),
+        )
 
-    @property
-    def viewing_url(self) -> str:
-        return f"{self.base_url}/home/local-job-progress/{str(self.job_uuid)}"
+    def generate_model_queues_table(self):
+        table = Table(show_header=False, box=box.SIMPLE)
+        table.add_column("Info", style="cyan")
+        table.add_column("Value", style="magenta")
+        # table.add_row("Bucket collection", str(self.jobs_runner.bucket_collection))
+        for model, bucket in self.jobs_runner.bucket_collection.items():
+            table.add_row(Text(model.model, style="bold blue"), "")
+            bucket_types = ["requests_bucket", "tokens_bucket"]
+            for bucket_type in bucket_types:
+                table.add_row(Text(" " + bucket_type, style="green"), "")
+                # table.add_row(
+                #     f"  Current level (capacity = {round(getattr(bucket, bucket_type).capacity, 3)})",
+                #     str(round(getattr(bucket, bucket_type).tokens, 3)),
+                # )
+                num_requests = getattr(bucket, bucket_type).num_requests
+                num_released = getattr(bucket, bucket_type).num_released
+                tokens_returned = getattr(bucket, bucket_type).tokens_returned
+                # table.add_row(
+                #     f"  Requested",
+                #     str(num_requests),
+                # )
+                # table.add_row(
+                #     f"  Completed",
+                #     str(num_released),
+                # )
+                table.add_row(
+                    "  Completed vs. Requested", f"{num_released} vs. {num_requests}"
+                )
+                table.add_row(
+                    "  Added tokens (from cache)",
+                    str(tokens_returned),
+                )
+                if bucket_type == "tokens_bucket":
+                    rate_name = "TPM"
+                else:
+                    rate_name = "RPM"
+                target_rate = round(getattr(bucket, bucket_type).target_rate, 1)
+                table.add_row(
+                    f"  Empirical {rate_name} (target = {target_rate})",
+                    str(round(getattr(bucket, bucket_type).get_throughput(), 0)),
+                )
+
+        return table
+
+    def generate_layout(self):
+        progress = self.create_progress_bar()
+        task_ids = []
+        for model in self.distinct_models:
+            task_id = progress.add_task(
+                f"[cyan]{model}...",
+                total=int(self.num_total_interviews / len(self.distinct_models)),
+            )
+            task_ids.append((model, task_id))
+
+        progress_height = min(5, 2 + len(self.distinct_models))
+        layout = Layout()
+
+        # Create the top row with only the progress panel
+        layout.split_column(
+            Layout(
+                Panel(
+                    progress,
+                    title="Interview Progress",
+                    border_style="cyan",
+                    box=box.ROUNDED,
+                ),
+                name="progress",
+                size=progress_height,  # Adjusted size
+            ),
+            Layout(name="bottom_row"),  # Adjusted size
+        )
 
-    @property
-    def update_url(self) -> str:
-        return f"{self.base_url}/api/v0/local-job/{str(self.job_uuid)}"
+        # Split the bottom row into two columns for metrics and model queues
+        layout["bottom_row"].split_row(
+            Layout(
+                Panel(
+                    self.generate_metrics_table(),
+                    title="Metrics",
+                    border_style="magenta",
+                    box=box.ROUNDED,
+                ),
+                name="metrics",
+            ),
+            Layout(
+                Panel(
+                    self.generate_model_queues_table(),
+                    title="Model Queues",
+                    border_style="yellow",
+                    box=box.ROUNDED,
+                ),
+                name="model_queues",
+            ),
+        )
 
-    def setup(self) -> None:
-        """
-        Creates a local job on Coop if one does not already exist.
-        """
+        return layout, progress, task_ids
 
-        headers = {"Content-Type": "application/json"}
+    def generate_metrics_table(self):
+        table = Table(show_header=True, header_style="bold magenta", box=box.SIMPLE)
+        table.add_column("Metric", style="cyan", no_wrap=True)
+        table.add_column("Value", justify="right")
 
-        if self.api_key:
-            headers["Authorization"] = f"Bearer {self.api_key}"
-        else:
-            headers["Authorization"] = f"Bearer None"
+        for stat_name in self.statistics:
+            pretty_name, value = list(self._compute_statistic(stat_name).items())[0]
+            # breakpoint()
+            table.add_row(pretty_name, value)
+        return table
 
-        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")
-
-        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:
-            # 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"
-
-            # Send the update
-            response = requests.patch(
-                self.update_url,
-                json=status_dict,
-                headers=headers,
-                timeout=1,
+    def update_progress(self, stop_event):
+        layout, progress, task_ids = self.generate_layout()
+
+        with Live(
+            layout, refresh_per_second=int(1 / self.refresh_rate), transient=True
+        ) as live:
+            while (
+                len(self.completed_interviews) < len(self.jobs_runner.total_interviews)
+                and not stop_event.is_set()
+            ):
+                completed_tasks = len(self.completed_interviews)
+                total_tasks = len(self.jobs_runner.total_interviews)
+
+                for model, task_id in task_ids:
+                    completed_tasks = len(self.completed_interview_by_model[model])
+                    progress.update(
+                        task_id,
+                        completed=completed_tasks,
+                        description=f"[cyan]Conducting interviews for {model}...",
+                    )
+
+                layout["metrics"].update(
+                    Panel(
+                        self.generate_metrics_table(),
+                        title="Metrics",
+                        border_style="magenta",
+                        box=box.ROUNDED,
+                    )
+                )
+                layout["model_queues"].update(
+                    Panel(
+                        self.generate_model_queues_table(),
+                        title="Final Model Queues",
+                        border_style="yellow",
+                        box=box.ROUNDED,
+                    )
+                )
+
+                time.sleep(self.refresh_rate)
+
+            # Final update
+            for model, task_id in task_ids:
+                completed_tasks = len(self.completed_interview_by_model[model])
+                progress.update(
+                    task_id,
+                    completed=completed_tasks,
+                    description=f"[cyan]Conducting interviews for {model}...",
+                )
+
+            layout["metrics"].update(
+                Panel(
+                    self.generate_metrics_table(),
+                    title="Final Metrics",
+                    border_style="magenta",
+                    box=box.ROUNDED,
+                )
             )
-            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. Otherwise, returns False.
-        """
-
-        if self.api_key is not None:
-            return True
-        else:
-            return False
+            live.update(layout)
+            time.sleep(1)  # Show final state for 1 second
 
 
 if __name__ == "__main__":

edsl/jobs/tasks/TaskHistory.py

@@ -8,12 +8,7 @@ from edsl.jobs.tasks.task_status_enum import TaskStatus
 
 
 class TaskHistory:
-    def __init__(
-        self,
-        interviews: List["Interview"],
-        include_traceback: bool = False,
-        max_interviews: int = 10,
-    ):
+    def __init__(self, interviews: List["Interview"], include_traceback: bool = False):
         """
         The structure of a TaskHistory exception
 
@@ -27,7 +22,6 @@ class TaskHistory:
         self.include_traceback = include_traceback
 
         self._interviews = {index: i for index, i in enumerate(self.total_interviews)}
-        self.max_interviews = max_interviews
 
     @classmethod
     def example(cls):
@@ -79,21 +73,19 @@ class TaskHistory:
         """Return a string representation of the TaskHistory."""
         return f"TaskHistory(interviews={self.total_interviews})."
 
-    def to_dict(self, add_edsl_version=True):
+    def to_dict(self):
         """Return the TaskHistory as a dictionary."""
-        d = {
-            "interviews": [
-                i.to_dict(add_edsl_version=add_edsl_version)
-                for i in self.total_interviews
-            ],
+        # return {
+        #     "exceptions": [
+        #         e.to_dict(include_traceback=self.include_traceback)
+        #         for e in self.exceptions
+        #     ],
+        #     "indices": self.indices,
+        # }
+        return {
+            "interviews": [i._to_dict() for i in self.total_interviews],
             "include_traceback": self.include_traceback,
         }
-        if add_edsl_version:
-            from edsl import __version__
-
-            d["edsl_version"] = __version__
-            d["edsl_class_name"] = "TaskHistory"
-        return d
 
     @classmethod
     def from_dict(cls, data: dict):
@@ -123,11 +115,10 @@ class TaskHistory:
 
     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
+        from edsl.utilities.utilities import data_to_html
 
-        return tabulate(data, headers=["keys", "values"], tablefmt="html")
+        newdata = self.to_dict()["exceptions"]
+        return data_to_html(newdata, replace_new_lines=True)
 
     def show_exceptions(self, tracebacks=False):
         """Print the exceptions."""
@@ -257,6 +248,8 @@ class TaskHistory:
             for question_name, exceptions in interview.exceptions.items():
                 for exception in exceptions:
                     exception_type = exception.exception.__class__.__name__
+                    # exception_type = exception["exception"]
+                    # breakpoint()
                     if exception_type in exceptions_by_type:
                         exceptions_by_type[exception_type] += 1
                     else:
@@ -343,9 +336,9 @@ class TaskHistory:
 
         env = Environment(loader=TemplateLoader("edsl", "templates/error_reporting"))
 
-        # Get current memory usage at this point
-
+        # Load and render a template
         template = env.get_template("base.html")
+        # rendered_template = template.render(your_data=your_data)
 
         # Render the template with data
         output = template.render(
@@ -359,7 +352,6 @@ class TaskHistory:
             exceptions_by_model=self.exceptions_by_model,
             exceptions_by_service=self.exceptions_by_service,
             models_used=models_used,
-            max_interviews=self.max_interviews,
         )
         return output
 
@@ -369,7 +361,7 @@ class TaskHistory:
         return_link=False,
         css=None,
         cta="Open Report in New Tab",
-        open_in_browser=False,
+        open_in_browser=True,
     ):
         """Return an HTML report."""