edsl 0.1.32__py3-none-any.whl → 0.1.33__py3-none-any.whl

edsl/jobs/runners/JobsRunnerStatus.py

@@ -0,0 +1,331 @@
+from __future__ import annotations
+
+import time
+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
+
+
+@dataclass
+class ModelInfo:
+    model_name: str
+    TPM_limit_k: float
+    RPM_limit_k: float
+    num_tasks_waiting: int
+    token_usage_info: dict
+
+
+@dataclass
+class ModelTokenUsageStats:
+    token_usage_type: str
+    details: List[dict]
+    cost: str
+
+
+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 = 0.25
+    ):
+        self.jobs_runner = jobs_runner
+        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.distinct_models = list(
+            set(i.model.model for i in self.jobs_runner.interviews)
+        )
+
+        self.completed_interview_by_model = defaultdict(list)
+
+    def add_completed_interview(self, result):
+        self.completed_interviews.append(result.interview_hash)
+
+        relevant_model = result.model.model
+        self.completed_interview_by_model[relevant_model].append(result.interview_hash)
+
+    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
+
+        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]()
+
+    def create_progress_bar(self):
+        return Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+            TextColumn("{task.completed}/{task.total}"),
+        )
+
+    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
+        )
+
+        # 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",
+            ),
+        )
+
+        return layout, progress, task_ids
+
+    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")
+
+        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
+
+    def update_progress(self):
+        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
+            ):
+                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,
+                )
+            )
+            live.update(layout)
+            time.sleep(1)  # Show final state for 1 second
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/tasks/QuestionTaskCreator.py

@@ -55,6 +55,7 @@ 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
@@ -87,10 +88,10 @@ class QuestionTaskCreator(UserList):
         """
         self.append(task)
 
-    def generate_task(self, debug: bool) -> asyncio.Task:
+    def generate_task(self) -> asyncio.Task:
         """Create a task that depends on the passed-in dependencies."""
         task = asyncio.create_task(
-            self._run_task_async(debug), name=self.question.question_name
+            self._run_task_async(), name=self.question.question_name
         )
         task.depends_on = [t.get_name() for t in self]
         return task
@@ -103,7 +104,7 @@ class QuestionTaskCreator(UserList):
         """Returns the token usage for the task.
 
         >>> qt = QuestionTaskCreator.example()
-        >>> answers = asyncio.run(qt._run_focal_task(debug=False))
+        >>> answers = asyncio.run(qt._run_focal_task())
         >>> qt.token_usage()
         {'cached_tokens': TokenUsage(from_cache=True, prompt_tokens=0, completion_tokens=0), 'new_tokens': TokenUsage(from_cache=False, prompt_tokens=0, completion_tokens=0)}
         """
@@ -111,15 +112,15 @@ class QuestionTaskCreator(UserList):
             cached_tokens=self.cached_token_usage, new_tokens=self.new_token_usage
         )
 
-    async def _run_focal_task(self, debug: bool) -> Answers:
+    async def _run_focal_task(self) -> Answers:
         """Run the focal task i.e., the question that we are interested in answering.
 
         It is only called after all the dependency tasks are completed.
 
         >>> qt = QuestionTaskCreator.example()
-        >>> answers = asyncio.run(qt._run_focal_task(debug=False))
-        >>> answers["answer"]
-        'Yo!'
+        >>> answers = asyncio.run(qt._run_focal_task())
+        >>> answers.answer
+        'This is an example answer'
         """
 
         requested_tokens = self.estimated_tokens()
@@ -132,19 +133,19 @@ class QuestionTaskCreator(UserList):
             self.waiting = True
             self.task_status = TaskStatus.WAITING_FOR_REQUEST_CAPACITY
 
-        await self.tokens_bucket.get_tokens(1)
+        await self.requests_bucket.get_tokens(1, cheat_bucket_capacity=True)
 
         self.task_status = TaskStatus.API_CALL_IN_PROGRESS
         try:
             results = await self.answer_question_func(
-                question=self.question, debug=debug, task=None  # self
+                question=self.question, task=None  # self
             )
             self.task_status = TaskStatus.SUCCESS
         except Exception as e:
             self.task_status = TaskStatus.FAILED
             raise e
 
-        if results.get("cache_used", False):
+        if results.cache_used:
             self.tokens_bucket.add_tokens(requested_tokens)
             self.requests_bucket.add_tokens(1)
             self.from_cache = True
@@ -155,17 +156,18 @@ class QuestionTaskCreator(UserList):
             self.tokens_bucket.turbo_mode_off()
             self.requests_bucket.turbo_mode_off()
 
-        _ = results.pop("cached_response", None)
+        # breakpoint()
+        # _ = results.pop("cached_response", None)
 
-        tracker = self.cached_token_usage if self.from_cache else self.new_token_usage
+        # tracker = self.cached_token_usage if self.from_cache else self.new_token_usage
 
         # TODO: This is hacky. The 'func' call should return an object that definitely has a 'usage' key.
-        usage = results.get("usage", {"prompt_tokens": 0, "completion_tokens": 0})
-        prompt_tokens = usage.get("prompt_tokens", 0)
-        completion_tokens = usage.get("completion_tokens", 0)
-        tracker.add_tokens(
-            prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
-        )
+        # usage = results.get("usage", {"prompt_tokens": 0, "completion_tokens": 0})
+        # prompt_tokens = usage.get("prompt_tokens", 0)
+        # completion_tokens = usage.get("completion_tokens", 0)
+        # tracker.add_tokens(
+        #    prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
+        # )
 
         return results
 
@@ -177,8 +179,13 @@ class QuestionTaskCreator(UserList):
 
         m = ModelBuckets.infinity_bucket()
 
-        async def answer_question_func(question, debug, task):
-            return {"answer": "Yo!"}
+        from collections import namedtuple
+
+        AnswerDict = namedtuple("AnswerDict", ["answer", "cache_used"])
+        answer = AnswerDict(answer="This is an example answer", cache_used=False)
+
+        async def answer_question_func(question, task):
+            return answer
 
         return cls(
             question=QuestionFreeText.example(),
@@ -188,7 +195,7 @@ class QuestionTaskCreator(UserList):
             iteration=0,
         )
 
-    async def _run_task_async(self, debug) -> None:
+    async def _run_task_async(self) -> None:
         """Run the task asynchronously, awaiting the tasks that must be completed before this one can be run.
 
         >>> qt1 = QuestionTaskCreator.example()
@@ -231,8 +238,6 @@ class QuestionTaskCreator(UserList):
                 if isinstance(result, Exception):
                     raise result
 
-            return await self._run_focal_task(debug)
-
         except asyncio.CancelledError:
             self.task_status = TaskStatus.CANCELLED
             raise
@@ -244,6 +249,8 @@ class QuestionTaskCreator(UserList):
                 f"Required tasks failed for {self.question.question_name}"
             ) from e
 
+        return await self._run_focal_task()
+
 
 if __name__ == "__main__":
     import doctest