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

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -3,25 +3,40 @@ import time
 import math
 import asyncio
 import functools
-import threading
 from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator
 from contextlib import contextmanager
 from collections import UserList
 
-from edsl.results.Results import Results
-from rich.live import Live
-from rich.console import Console
-
 from edsl import shared_globals
 from edsl.jobs.interviews.Interview import Interview
-from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
-
+from edsl.jobs.runners.JobsRunnerStatusMixin import JobsRunnerStatusMixin
 from edsl.jobs.tasks.TaskHistory import TaskHistory
 from edsl.jobs.buckets.BucketCollection import BucketCollection
 from edsl.utilities.decorators import jupyter_nb_handler
 from edsl.data.Cache import Cache
 from edsl.results.Result import Result
 from edsl.results.Results import Results
+from edsl.jobs.FailedQuestion import FailedQuestion
+
+
+def cache_with_timeout(timeout):
+    """ "Used to keep the generate table from being run too frequetly."""
+
+    def decorator(func):
+        cached_result = {}
+        last_computation_time = [0]  # Using list to store mutable value
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            current_time = time.time()
+            if (current_time - last_computation_time[0]) >= timeout:
+                cached_result["value"] = func(*args, **kwargs)
+                last_computation_time[0] = current_time
+            return cached_result["value"]
+
+        return wrapper
+
+    return decorator
 
 
 class StatusTracker(UserList):
@@ -33,7 +48,7 @@ class StatusTracker(UserList):
         return print(f"Completed: {len(self.data)} of {self.total_tasks}", end="\r")
 
 
-class JobsRunnerAsyncio:
+class JobsRunnerAsyncio(JobsRunnerStatusMixin):
     """A class for running a collection of interviews asynchronously.
 
     It gets instaniated from a Jobs object.
@@ -42,12 +57,11 @@ class JobsRunnerAsyncio:
 
     def __init__(self, jobs: "Jobs"):
         self.jobs = jobs
+        # this creates the interviews, which can take a while
         self.interviews: List["Interview"] = jobs.interviews()
         self.bucket_collection: "BucketCollection" = jobs.bucket_collection
         self.total_interviews: List["Interview"] = []
 
-        # self.jobs_runner_status = JobsRunnerStatus(self, n=1)
-
     async def run_async_generator(
         self,
         cache: "Cache",
@@ -65,7 +79,6 @@ class JobsRunnerAsyncio:
         :param stop_on_exception: Whether to stop the interview if an exception is raised
         :param sidecar_model: a language model to use in addition to the interview's model
         :param total_interviews: A list of interviews to run can be provided instead.
-        :param raise_validation_errors: Whether to raise validation errors
         """
         tasks = []
         if total_interviews:  # was already passed in total interviews
@@ -75,6 +88,8 @@ class JobsRunnerAsyncio:
                 self._populate_total_interviews(n=n)
             )  # Populate self.total_interviews before creating tasks
 
+        # print("Interviews created")
+
         for interview in self.total_interviews:
             interviewing_task = self._build_interview_task(
                 interview=interview,
@@ -84,9 +99,11 @@ class JobsRunnerAsyncio:
             )
             tasks.append(asyncio.create_task(interviewing_task))
 
+        # print("Tasks created")
+
         for task in asyncio.as_completed(tasks):
+            # print(f"Task {task} completed")
             result = await task
-            self.jobs_runner_status.add_completed_interview(result)
             yield result
 
     def _populate_total_interviews(
@@ -105,8 +122,6 @@ class JobsRunnerAsyncio:
                     yield interview
 
     async def run_async(self, cache: Optional["Cache"] = None, n: int = 1) -> Results:
-        """Used for some other modules that have a non-standard way of running interviews."""
-        self.jobs_runner_status = JobsRunnerStatus(self, n=n)
         self.cache = Cache() if cache is None else cache
         data = []
         async for result in self.run_async_generator(cache=self.cache, n=n):
@@ -142,6 +157,12 @@ class JobsRunnerAsyncio:
             raise_validation_errors=raise_validation_errors,
         )
 
+        # answer_key_names = {
+        #     k
+        #     for k in set(answer.keys())
+        #     if not k.endswith("_comment") and not k.endswith("_generated_tokens")
+        # }
+
         question_results = {}
         for result in valid_results:
             question_results[result.question_name] = result
@@ -153,13 +174,24 @@ class JobsRunnerAsyncio:
             for k in answer_key_names
         }
         comments_dict = {
-            k + "_comment": question_results[k].comment for k in answer_key_names
+            "k" + "_comment": question_results[k].comment for k in answer_key_names
         }
 
         # we should have a valid result for each question
         answer_dict = {k: answer[k] for k in answer_key_names}
         assert len(valid_results) == len(answer_key_names)
 
+        # breakpoint()
+        # generated_tokens_dict = {
+        #     k + "_generated_tokens": v.generated_tokens
+        #     for k, v in zip(answer_key_names, valid_results)
+        # }
+
+        # comments_dict = {
+        #    k + "_comment": v.comment for k, v in zip(answer_key_names, valid_results)
+        # }
+        # breakpoint()
+
         # TODO: move this down into Interview
         question_name_to_prompts = dict({})
         for result in valid_results:
@@ -171,19 +203,19 @@ class JobsRunnerAsyncio:
 
         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"]
+            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
+            raw_model_results_dictionary[question_name + "_raw_model_response"] = (
+                result.raw_model_response
+            )
             raw_model_results_dictionary[question_name + "_cost"] = result.cost
             one_use_buys = (
                 "NA"
@@ -194,6 +226,7 @@ class JobsRunnerAsyncio:
             )
             raw_model_results_dictionary[question_name + "_one_usd_buys"] = one_use_buys
 
+        # breakpoint()
         result = Result(
             agent=interview.agent,
             scenario=interview.scenario,
@@ -214,62 +247,6 @@ class JobsRunnerAsyncio:
     def elapsed_time(self):
         return time.monotonic() - self.start_time
 
-    def process_results(
-        self, raw_results: Results, cache: Cache, print_exceptions: bool
-    ):
-        interview_lookup = {
-            hash(interview): index
-            for index, interview in enumerate(self.total_interviews)
-        }
-        interview_hashes = list(interview_lookup.keys())
-
-        results = Results(
-            survey=self.jobs.survey,
-            data=sorted(
-                raw_results, key=lambda x: interview_hashes.index(x.interview_hash)
-            ),
-        )
-        results.cache = cache
-        results.task_history = TaskHistory(
-            self.total_interviews, include_traceback=False
-        )
-        results.has_unfixed_exceptions = results.task_history.has_unfixed_exceptions
-        results.bucket_collection = self.bucket_collection
-
-        if results.has_unfixed_exceptions and print_exceptions:
-            from edsl.scenarios.FileStore import HTMLFileStore
-            from edsl.config import CONFIG
-            from edsl.coop.coop import Coop
-
-            msg = f"Exceptions were raised in {len(results.task_history.indices)} out of {len(self.total_interviews)} interviews.\n"
-
-            if len(results.task_history.indices) > 5:
-                msg += f"Exceptions were raised in the following interviews: {results.task_history.indices}.\n"
-
-            print(msg)
-            # this is where exceptions are opening up
-            filepath = results.task_history.html(
-                cta="Open report to see details.",
-                open_in_browser=True,
-                return_link=True,
-            )
-
-            try:
-                coop = Coop()
-                user_edsl_settings = coop.edsl_settings
-                remote_logging = user_edsl_settings["remote_logging"]
-            except Exception as e:
-                print(e)
-                remote_logging = False
-            if 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")
-
-        return results
-
     @jupyter_nb_handler
     async def run(
         self,
@@ -282,16 +259,24 @@ class JobsRunnerAsyncio:
         raise_validation_errors: bool = False,
     ) -> "Coroutine":
         """Runs a collection of interviews, handling both async and sync contexts."""
+        from rich.console import Console
 
+        console = Console()
         self.results = []
         self.start_time = time.monotonic()
         self.completed = False
         self.cache = cache
         self.sidecar_model = sidecar_model
 
-        self.jobs_runner_status = JobsRunnerStatus(self, n=n)
+        from edsl.results.Results import Results
+        from rich.live import Live
+        from rich.console import Console
 
-        async def process_results(cache):
+        @cache_with_timeout(1)
+        def generate_table():
+            return self.status_table(self.results, self.elapsed_time)
+
+        async def process_results(cache, progress_bar_context=None):
             """Processes results from interviews."""
             async for result in self.run_async_generator(
                 n=n,
@@ -301,22 +286,112 @@ class JobsRunnerAsyncio:
                 raise_validation_errors=raise_validation_errors,
             ):
                 self.results.append(result)
+                if progress_bar_context:
+                    progress_bar_context.update(generate_table())
             self.completed = True
 
-        def run_progress_bar():
-            """Runs the progress bar in a separate thread."""
-            self.jobs_runner_status.update_progress()
+        async def update_progress_bar(progress_bar_context):
+            """Updates the progress bar at fixed intervals."""
+            if progress_bar_context is None:
+                return
+
+            while True:
+                progress_bar_context.update(generate_table())
+                await asyncio.sleep(0.1)  # Update interval
+                if self.completed:
+                    break
+
+        @contextmanager
+        def conditional_context(condition, context_manager):
+            if condition:
+                with context_manager as cm:
+                    yield cm
+            else:
+                yield
+
+        with conditional_context(
+            progress_bar, Live(generate_table(), console=console, refresh_per_second=1)
+        ) as progress_bar_context:
+            with cache as c:
+                progress_task = asyncio.create_task(
+                    update_progress_bar(progress_bar_context)
+                )
+
+                try:
+                    await asyncio.gather(
+                        progress_task,
+                        process_results(
+                            cache=c, progress_bar_context=progress_bar_context
+                        ),
+                    )
+                except asyncio.CancelledError:
+                    pass
+                finally:
+                    progress_task.cancel()  # Cancel the progress_task when process_results is done
+                    await progress_task
+
+                    await asyncio.sleep(1)  # short delay to show the final status
+
+                    if progress_bar_context:
+                        progress_bar_context.update(generate_table())
+
+        # puts results in the same order as the total interviews
+        interview_lookup = {
+            hash(interview): index
+            for index, interview in enumerate(self.total_interviews)
+        }
+        interview_hashes = list(interview_lookup.keys())
+        self.results = sorted(
+            self.results, key=lambda x: interview_hashes.index(x.interview_hash)
+        )
+
+        results = Results(survey=self.jobs.survey, data=self.results)
+        task_history = TaskHistory(self.total_interviews, include_traceback=False)
+        results.task_history = task_history
 
-        if progress_bar:
-            progress_thread = threading.Thread(target=run_progress_bar)
-            progress_thread.start()
+        results.failed_questions = {}
+        results.has_exceptions = task_history.has_exceptions
 
-        with cache as c:
-            await process_results(cache=c)
+        # breakpoint()
+        results.bucket_collection = self.bucket_collection
 
-        if progress_bar:
-            progress_thread.join()
+        if results.has_exceptions:
+            # put the failed interviews in the results object as a list
+            failed_interviews = [
+                interview.duplicate(
+                    iteration=interview.iteration, cache=interview.cache
+                )
+                for interview in self.total_interviews
+                if interview.has_exceptions
+            ]
 
-        return self.process_results(
-            raw_results=self.results, cache=cache, print_exceptions=print_exceptions
-        )
+            failed_questions = {}
+            for interview in self.total_interviews:
+                if interview.has_exceptions:
+                    index = interview_lookup[hash(interview)]
+                    failed_questions[index] = interview.failed_questions
+
+            results.failed_questions = failed_questions
+
+            from edsl.jobs.Jobs import Jobs
+
+            results.failed_jobs = Jobs.from_interviews(
+                [interview for interview in failed_interviews]
+            )
+            if print_exceptions:
+                msg = f"Exceptions were raised in {len(results.task_history.indices)} out of {len(self.total_interviews)} interviews.\n"
+
+                if len(results.task_history.indices) > 5:
+                    msg += f"Exceptions were raised in the following interviews: {results.task_history.indices}.\n"
+
+                shared_globals["edsl_runner_exceptions"] = task_history
+                print(msg)
+                # this is where exceptions are opening up
+                task_history.html(
+                    cta="Open report to see details.", open_in_browser=True
+                )
+                print(
+                    "Also see: https://docs.expectedparrot.com/en/latest/exceptions.html"
+                )
+
+        return results