edsl 0.1.31.dev3__py3-none-any.whl → 0.1.32__py3-none-any.whl

edsl/jobs/runners/JobsRunnerStatusData.py

@@ -1,237 +0,0 @@
-import asyncio
-from enum import Enum
-from typing import Literal, List, Type, DefaultDict
-from collections import UserDict, defaultdict
-
-from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
-from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
-
-# from edsl.enums import pricing, TokenPricing
-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,
-)
-
-
-class JobsRunnerStatusData:
-    def status_dict(
-        self, interviews: List[Type["Interview"]]
-    ) -> List[Type[InterviewStatusDictionary]]:
-        """
-        >>> from edsl.jobs.interviews.Interview import Interview
-        >>> interviews = [Interview.example()]
-        >>> JobsRunnerStatusData().status_dict(interviews)
-        [InterviewStatusDictionary({<TaskStatus.NOT_STARTED: 1>: 0, <TaskStatus.WAITING_FOR_DEPENDENCIES: 2>: 0, <TaskStatus.CANCELLED: 3>: 0, <TaskStatus.PARENT_FAILED: 4>: 0, <TaskStatus.WAITING_FOR_REQUEST_CAPACITY: 5>: 0, <TaskStatus.WAITING_FOR_TOKEN_CAPACITY: 6>: 0, <TaskStatus.API_CALL_IN_PROGRESS: 7>: 0, <TaskStatus.SUCCESS: 8>: 0, <TaskStatus.FAILED: 9>: 0, 'number_from_cache': 0})]
-        """
-        status = []
-        for interview in interviews:
-            status.append(interview.interview_status)
-
-        return status
-
-    def status_counts(self, interviews: List[Type["Interview"]]):
-        """
-        Takes a collection of interviews and returns a dictionary of the counts of each status.
-
-        :param interviews: a collection of interviews.
-
-        This creates a dictionary of the counts of each status in the collection of interviews.
-
-        >>> from edsl.jobs.interviews.Interview import Interview
-        >>> interviews = [Interview.example() for _ in range(100)]
-        >>> jd = JobsRunnerStatusData()
-        >>> jd.status_counts(interviews)
-        dict_values([InterviewStatusDictionary({<TaskStatus.NOT_STARTED: 1>: 0, <TaskStatus.WAITING_FOR_DEPENDENCIES: 2>: 0, <TaskStatus.CANCELLED: 3>: 0, <TaskStatus.PARENT_FAILED: 4>: 0, <TaskStatus.WAITING_FOR_REQUEST_CAPACITY: 5>: 0, <TaskStatus.WAITING_FOR_TOKEN_CAPACITY: 6>: 0, <TaskStatus.API_CALL_IN_PROGRESS: 7>: 0, <TaskStatus.SUCCESS: 8>: 0, <TaskStatus.FAILED: 9>: 0, 'number_from_cache': 0})])
-        >>> len(jd.status_counts(interviews))
-        1
-        """
-        model_to_status = defaultdict(InterviewStatusDictionary)
-
-        for interview in interviews:
-            model = interview.model  # get the model for the interview
-            model_to_status[
-                model
-            ] += (
-                interview.interview_status
-            )  # InterviewStatusDictionary objects can be added together
-
-        return (
-            model_to_status.values()
-        )  # return the values of the dictionary, which is a list of dictionaries
-
-    def generate_status_summary(
-        self,
-        completed_tasks: List[Type[asyncio.Task]],
-        elapsed_time: float,
-        interviews: List[Type["Interview"]],
-    ) -> InterviewStatisticsCollection:
-        """Generate a summary of the status of the job runner.
-
-        :param completed_tasks: list of completed tasks
-        :param elapsed_time: time elapsed since the start of the job
-        :param interviews: list of interviews to be conducted
-
-        >>> from edsl.jobs.interviews.Interview import Interview
-        >>> interviews = [Interview.example()]
-        >>> completed_tasks = []
-        >>> elapsed_time = 0
-        >>> JobsRunnerStatusData().generate_status_summary(completed_tasks, elapsed_time, interviews)
-        {'Elapsed time': '0.0 sec.', 'Total interviews requested': '1 ', 'Completed interviews': '0 ', 'Percent complete': '0 %', 'Average time per interview': 'NA', 'Task remaining': '1 ', 'Estimated time remaining': 'NA', 'model_queues': [{'model_name': '...', 'TPM_limit_k': ..., 'RPM_limit_k': ..., 'num_tasks_waiting': 0, 'token_usage_info': [{'cache_status': 'new_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}, {'cache_status': 'cached_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}]}]}
-        """
-
-        models_to_tokens = defaultdict(InterviewTokenUsage)
-        model_to_status = defaultdict(InterviewStatusDictionary)
-
-        waiting_dict = defaultdict(int)
-
-        interview_statistics = InterviewStatisticsCollection()
-
-        for interview in interviews:
-            model = interview.model
-            models_to_tokens[model] += interview.token_usage
-            model_to_status[model] += interview.interview_status
-            waiting_dict[model] += interview.interview_status.waiting
-
-        interview_statistics.add_stat(
-            InterviewStatistic(
-                "elapsed_time", value=elapsed_time, digits=1, units="sec."
-            )
-        )
-        interview_statistics.add_stat(
-            InterviewStatistic(
-                "total_interviews_requested", value=len(interviews), units=""
-            )
-        )
-        interview_statistics.add_stat(
-            InterviewStatistic(
-                "completed_interviews", value=len(completed_tasks), units=""
-            )
-        )
-        interview_statistics.add_stat(
-            InterviewStatistic(
-                "percent_complete",
-                value=(
-                    len(completed_tasks) / len(interviews) * 100
-                    if len(interviews) > 0
-                    else "NA"
-                ),
-                digits=0,
-                units="%",
-            )
-        )
-        interview_statistics.add_stat(
-            InterviewStatistic(
-                "average_time_per_interview",
-                value=elapsed_time / len(completed_tasks) if completed_tasks else "NA",
-                digits=1,
-                units="sec.",
-            )
-        )
-        interview_statistics.add_stat(
-            InterviewStatistic(
-                "task_remaining", value=len(interviews) - len(completed_tasks), units=""
-            )
-        )
-        number_remaining = len(interviews) - len(completed_tasks)
-        time_per_task = (
-            elapsed_time / len(completed_tasks) if len(completed_tasks) > 0 else "NA"
-        )
-        estimated_time_remaining = (
-            number_remaining * time_per_task if time_per_task != "NA" else "NA"
-        )
-
-        interview_statistics.add_stat(
-            InterviewStatistic(
-                "estimated_time_remaining",
-                value=estimated_time_remaining,
-                digits=1,
-                units="sec.",
-            )
-        )
-        model_queues_info = []
-        for model, num_waiting in waiting_dict.items():
-            model_info = self._get_model_info(model, num_waiting, models_to_tokens)
-            model_queues_info.append(model_info)
-
-        interview_statistics["model_queues"] = model_queues_info
-
-        return interview_statistics
-
-    def _get_model_info(
-        self,
-        model: str,
-        num_waiting: int,
-        models_to_tokens: InterviewTokenUsageMapping,
-    ) -> dict:
-        """Get the status of a model.
-
-        >>> from edsl.jobs.interviews.Interview import Interview
-        >>> interviews = [Interview.example()]
-        >>> models_to_tokens = defaultdict(InterviewTokenUsage)
-        >>> model = interviews[0].model
-        >>> num_waiting = 0
-        >>> JobsRunnerStatusData()._get_model_info(model, num_waiting, models_to_tokens)
-        {'model_name': 'gpt-4-1106-preview', 'TPM_limit_k': ..., 'RPM_limit_k': ..., 'num_tasks_waiting': 0, 'token_usage_info': [{'cache_status': 'new_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}, {'cache_status': 'cached_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}]}
-        """
-
-        prices = get_token_pricing(model.model)
-
-        model_info = {
-            "model_name": model.model,
-            "TPM_limit_k": model.TPM / 1000,
-            "RPM_limit_k": model.RPM / 1000,
-            "num_tasks_waiting": num_waiting,
-            "token_usage_info": [],
-        }
-
-        token_usage_types = ["new_token_usage", "cached_token_usage"]
-        for token_usage_type in token_usage_types:
-            cache_info = self._get_token_usage_info(
-                token_usage_type, models_to_tokens, model, prices
-            )
-            model_info["token_usage_info"].append(cache_info)
-
-        return model_info
-
-    def _get_token_usage_info(
-        self,
-        cache_status: Literal["new_token_usage", "cached_token_usage"],
-        models_to_tokens: InterviewTokenUsageMapping,
-        model: str,
-        prices: "TokenPricing",
-    ) -> dict:
-        """Get the token usage info for a model.
-
-        >>> from edsl.jobs.interviews.Interview import Interview
-        >>> interviews = [Interview.example()]
-        >>> models_to_tokens = defaultdict(InterviewTokenUsage)
-        >>> model = interviews[0].model
-        >>> prices = get_token_pricing(model.model)
-        >>> cache_status = "new_token_usage"
-        >>> JobsRunnerStatusData()._get_token_usage_info(cache_status, models_to_tokens, model, prices)
-        {'cache_status': 'new_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}
-
-        """
-        cache_info = {"cache_status": cache_status, "details": []}
-        token_usage = getattr(models_to_tokens[model], cache_status)
-        for token_type in ["prompt_tokens", "completion_tokens"]:
-            tokens = getattr(token_usage, token_type)
-            cache_info["details"].append(
-                {
-                    "type": token_type,
-                    "tokens": tokens,
-                }
-            )
-        cache_info["cost"] = f"${token_usage.cost(prices):.5f}"
-        return cache_info
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/runners/JobsRunnerStatusMixin.py

@@ -1,27 +1,272 @@
 from __future__ import annotations
 from typing import List, DefaultDict
 import asyncio
+from typing import Type
+from collections import defaultdict
+
+from typing import Literal, List, Type, DefaultDict
+from collections import UserDict, defaultdict
+
+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
+
+
+# return {"cache_status": token_usage_type, "details": details, "cost": f"${token_usage.cost(prices):.5f}"}
+
+from dataclasses import dataclass, asdict
 
 from rich.text import Text
 from rich.box import SIMPLE
 from rich.table import Table
 
-from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
-from edsl.jobs.runners.JobsRunnerStatusData import JobsRunnerStatusData
 
-InterviewTokenUsageMapping = DefaultDict[str, 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 JobsRunnerStatusMixin:
+    # @staticmethod
+    # def status_dict(interviews: List[Type["Interview"]]) -> List[Type[InterviewStatusDictionary]]:
+    #     """
+    #     >>> from edsl.jobs.interviews.Interview import Interview
+    #     >>> interviews = [Interview.example()]
+    #     >>> JobsRunnerStatusMixin().status_dict(interviews)
+    #     [InterviewStatusDictionary({<TaskStatus.NOT_STARTED: 1>: 0, <TaskStatus.WAITING_FOR_DEPENDENCIES: 2>: 0, <TaskStatus.CANCELLED: 3>: 0, <TaskStatus.PARENT_FAILED: 4>: 0, <TaskStatus.WAITING_FOR_REQUEST_CAPACITY: 5>: 0, <TaskStatus.WAITING_FOR_TOKEN_CAPACITY: 6>: 0, <TaskStatus.API_CALL_IN_PROGRESS: 7>: 0, <TaskStatus.SUCCESS: 8>: 0, <TaskStatus.FAILED: 9>: 0, 'number_from_cache': 0})]
+    #     """
+    #     return [interview.interview_status for interview in interviews]
 
+    def _compute_statistic(stat_name: str, completed_tasks, elapsed_time, 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 "NA"
+                ),
+                digits=0,
+                units="%",
+            ),
+            "average_time_per_interview": lambda: InterviewStatistic(
+                "average_time_per_interview",
+                value=elapsed_time / len(completed_tasks) if completed_tasks else "NA",
+                digits=1,
+                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 "NA"
+                ),
+                digits=1,
+                units="sec.",
+            ),
+        }
+        if stat_name not in stat_definitions:
+            raise ValueError(
+                f"Invalid stat_name: {stat_name}. The valid stat_names are: {list(stat_definitions.keys())}"
+            )
+        return stat_definitions[stat_name]()
 
-class JobsRunnerStatusPresentation:
     @staticmethod
-    def display_status_table(status_summary):
-        table = Table(
-            title="Job Status",
-            show_header=True,
-            header_style="bold magenta",
-            box=SIMPLE,
+    def _job_level_info(
+        completed_tasks: List[Type[asyncio.Task]],
+        elapsed_time: float,
+        interviews: List[Type["Interview"]],
+    ) -> InterviewStatisticsCollection:
+        interview_statistics = InterviewStatisticsCollection()
+
+        default_statistics = [
+            "elapsed_time",
+            "total_interviews_requested",
+            "completed_interviews",
+            "percent_complete",
+            "average_time_per_interview",
+            "task_remaining",
+            "estimated_time_remaining",
+        ]
+        for stat_name in default_statistics:
+            interview_statistics.add_stat(
+                JobsRunnerStatusMixin._compute_statistic(
+                    stat_name, completed_tasks, elapsed_time, interviews
+                )
+            )
+
+        return interview_statistics
+
+    @staticmethod
+    def _get_model_queues_info(interviews):
+        models_to_tokens = defaultdict(InterviewTokenUsage)
+        model_to_status = defaultdict(InterviewStatusDictionary)
+        waiting_dict = defaultdict(int)
+
+        for interview in interviews:
+            models_to_tokens[interview.model] += interview.token_usage
+            model_to_status[interview.model] += interview.interview_status
+            waiting_dict[interview.model] += interview.interview_status.waiting
+
+        for model, num_waiting in waiting_dict.items():
+            yield JobsRunnerStatusMixin._get_model_info(
+                model, num_waiting, models_to_tokens
+            )
+
+    @staticmethod
+    def generate_status_summary(
+        completed_tasks: List[Type[asyncio.Task]],
+        elapsed_time: float,
+        interviews: List[Type["Interview"]],
+        include_model_queues=False,
+    ) -> InterviewStatisticsCollection:
+        """Generate a summary of the status of the job runner.
+
+        :param completed_tasks: list of completed tasks
+        :param elapsed_time: time elapsed since the start of the job
+        :param interviews: list of interviews to be conducted
+
+        >>> from edsl.jobs.interviews.Interview import Interview
+        >>> interviews = [Interview.example()]
+        >>> completed_tasks = []
+        >>> elapsed_time = 0
+        >>> JobsRunnerStatusMixin().generate_status_summary(completed_tasks, elapsed_time, interviews)
+        {'Elapsed time': '0.0 sec.', 'Total interviews requested': '1 ', 'Completed interviews': '0 ', 'Percent complete': '0 %', 'Average time per interview': 'NA', 'Task remaining': '1 ', 'Estimated time remaining': 'NA'}
+        """
+
+        interview_status_summary: InterviewStatisticsCollection = (
+            JobsRunnerStatusMixin._job_level_info(
+                completed_tasks=completed_tasks,
+                elapsed_time=elapsed_time,
+                interviews=interviews,
+            )
         )
-        table.max_width = 100
+        if include_model_queues:
+            interview_status_summary.model_queues = list(
+                JobsRunnerStatusMixin._get_model_queues_info(interviews)
+            )
+        else:
+            interview_status_summary.model_queues = None
+
+        return interview_status_summary
+
+    @staticmethod
+    def _get_model_info(
+        model: str,
+        num_waiting: int,
+        models_to_tokens: InterviewTokenUsageMapping,
+    ) -> dict:
+        """Get the status of a model.
+
+        :param model: the model name
+        :param num_waiting: the number of tasks waiting for capacity
+        :param models_to_tokens: a mapping of models to token usage
+
+        >>> from edsl.jobs.interviews.Interview import Interview
+        >>> interviews = [Interview.example()]
+        >>> models_to_tokens = defaultdict(InterviewTokenUsage)
+        >>> model = interviews[0].model
+        >>> num_waiting = 0
+        >>> JobsRunnerStatusMixin()._get_model_info(model, num_waiting, models_to_tokens)
+        ModelInfo(model_name='gpt-4-1106-preview', TPM_limit_k=480.0, RPM_limit_k=4.0, num_tasks_waiting=0, token_usage_info=[ModelTokenUsageStats(token_usage_type='new_token_usage', details=[{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], cost='$0.00000'), ModelTokenUsageStats(token_usage_type='cached_token_usage', details=[{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], cost='$0.00000')])
+        """
+
+        ## TODO: This should probably be a coop method
+        prices = get_token_pricing(model.model)
+
+        token_usage_info = []
+        for token_usage_type in ["new_token_usage", "cached_token_usage"]:
+            token_usage_info.append(
+                JobsRunnerStatusMixin._get_token_usage_info(
+                    token_usage_type, models_to_tokens, model, prices
+                )
+            )
+
+        return ModelInfo(
+            **{
+                "model_name": model.model,
+                "TPM_limit_k": model.TPM / 1000,
+                "RPM_limit_k": model.RPM / 1000,
+                "num_tasks_waiting": num_waiting,
+                "token_usage_info": token_usage_info,
+            }
+        )
+
+    @staticmethod
+    def _get_token_usage_info(
+        token_usage_type: Literal["new_token_usage", "cached_token_usage"],
+        models_to_tokens: InterviewTokenUsageMapping,
+        model: str,
+        prices: "TokenPricing",
+    ) -> ModelTokenUsageStats:
+        """Get the token usage info for a model.
+
+        >>> from edsl.jobs.interviews.Interview import Interview
+        >>> interviews = [Interview.example()]
+        >>> models_to_tokens = defaultdict(InterviewTokenUsage)
+        >>> model = interviews[0].model
+        >>> prices = get_token_pricing(model.model)
+        >>> cache_status = "new_token_usage"
+        >>> JobsRunnerStatusMixin()._get_token_usage_info(cache_status, models_to_tokens, model, prices)
+        ModelTokenUsageStats(token_usage_type='new_token_usage', details=[{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], cost='$0.00000')
+
+        """
+        all_token_usage: InterviewTokenUsage = models_to_tokens[model]
+        token_usage: TokenUsage = getattr(all_token_usage, token_usage_type)
+
+        details = [
+            {"type": token_type, "tokens": getattr(token_usage, token_type)}
+            for token_type in ["prompt_tokens", "completion_tokens"]
+        ]
+
+        return ModelTokenUsageStats(
+            token_usage_type=token_usage_type,
+            details=details,
+            cost=f"${token_usage.cost(prices):.5f}",
+        )
+
+    @staticmethod
+    def _add_statistics_to_table(table, status_summary):
         table.add_column("Statistic", style="dim", no_wrap=True, width=50)
         table.add_column("Value", width=10)
 
@@ -29,49 +274,60 @@ class JobsRunnerStatusPresentation:
             if key != "model_queues":
                 table.add_row(key, value)
 
+    @staticmethod
+    def display_status_table(status_summary: InterviewStatisticsCollection) -> "Table":
+        table = Table(
+            title="Job Status",
+            show_header=True,
+            header_style="bold magenta",
+            box=SIMPLE,
+        )
+
+        ### Job-level statistics
+        JobsRunnerStatusMixin._add_statistics_to_table(table, status_summary)
+
+        ## Model-level statistics
         spacing = " "
-        if "model_queues" in status_summary:
+
+        if status_summary.model_queues is not None:
             table.add_row(Text("Model Queues", style="bold red"), "")
-            for model_info in status_summary["model_queues"]:
-                model_name = model_info["model_name"]
-                tpm = "TPM (k)=" + str(model_info["TPM_limit_k"])
-                rpm = "RPM (k)=" + str(model_info["RPM_limit_k"])
+            for model_info in status_summary.model_queues:
+                model_name = model_info.model_name
+                tpm = f"TPM (k)={model_info.TPM_limit_k}"
+                rpm = f"RPM (k)= {model_info.RPM_limit_k}"
                 pretty_model_name = model_name + ";" + tpm + ";" + rpm
                 table.add_row(Text(pretty_model_name, style="blue"), "")
                 table.add_row(
                     "Number question tasks waiting for capacity",
-                    str(model_info["num_tasks_waiting"]),
+                    str(model_info.num_tasks_waiting),
                 )
                 # Token usage and cost info
-                for cache_info in model_info["token_usage_info"]:
-                    cache_status = cache_info["cache_status"]
+                for token_usage_info in model_info.token_usage_info:
+                    token_usage_type = token_usage_info.token_usage_type
                     table.add_row(
-                        Text(spacing + cache_status.replace("_", " "), style="bold"), ""
+                        Text(
+                            spacing + token_usage_type.replace("_", " "), style="bold"
+                        ),
+                        "",
                     )
-                    for detail in cache_info["details"]:
+                    for detail in token_usage_info.details:
                         token_type = detail["type"]
                         tokens = detail["tokens"]
-                        #                 cost = detail["cost"]
                         table.add_row(spacing + f"{token_type}", f"{tokens:,}")
-                    table.add_row(spacing + "cost", cache_info["cost"])
+                    # table.add_row(spacing + "cost", cache_info["cost"])
 
         return table
 
-
-class JobsRunnerStatusMixin(JobsRunnerStatusData, JobsRunnerStatusPresentation):
-    def status_data(self, completed_tasks: List[asyncio.Task], elapsed_time: float):
-        # return self.generate_status_summary(
-        #     completed_tasks=completed_tasks,
-        #     elapsed_time=elapsed_time,
-        #     interviews=self.total_interviews).rawplt.figure(figsize=(10, 6))
-
-        # return self.full_status(self.total_interviews)
-        return None
-
     def status_table(self, completed_tasks: List[asyncio.Task], elapsed_time: float):
-        summary_data = self.generate_status_summary(
+        summary_data = JobsRunnerStatusMixin.generate_status_summary(
             completed_tasks=completed_tasks,
             elapsed_time=elapsed_time,
             interviews=self.total_interviews,
         )
         return self.display_status_table(summary_data)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/tasks/TaskCreators.py

@@ -17,7 +17,13 @@ class TaskCreators(UserDict):
 
     @property
     def token_usage(self) -> InterviewTokenUsage:
-        """Determines how many tokens were used for the interview."""
+        """Determines how many tokens were used for the interview.
+
+        This is iterates through all tasks that make up an interview.
+        For each task, it determines how many tokens were used and whether they were cached or new.
+        It then sums the total number of cached and new tokens used for the interview.
+
+        """
         cached_tokens = TokenUsage(from_cache=True)
         new_tokens = TokenUsage(from_cache=False)
         for task_creator in self.values():
@@ -28,7 +34,7 @@ class TaskCreators(UserDict):
             new_token_usage=new_tokens, cached_token_usage=cached_tokens
         )
 
-    def print(self):
+    def print(self) -> None:
         from rich import print
 
         print({task.get_name(): task.task_status for task in self.values()})