edsl 0.1.37.dev5__py3-none-any.whl → 0.1.38__py3-none-any.whl

edsl/jobs/interviews/Interview.py

@@ -159,7 +159,7 @@ class Interview:
         return self.task_creators.interview_status
 
     # region: Serialization
-    def _to_dict(self, include_exceptions=True) -> dict[str, Any]:
+    def to_dict(self, include_exceptions=True, add_edsl_version=True) -> dict[str, Any]:
         """Return a dictionary representation of the Interview instance.
         This is just for hashing purposes.
 
@@ -168,10 +168,10 @@ class Interview:
         1217840301076717434
         """
         d = {
-            "agent": self.agent._to_dict(),
-            "survey": self.survey._to_dict(),
-            "scenario": self.scenario._to_dict(),
-            "model": self.model._to_dict(),
+            "agent": self.agent.to_dict(add_edsl_version=add_edsl_version),
+            "survey": self.survey.to_dict(add_edsl_version=add_edsl_version),
+            "scenario": self.scenario.to_dict(add_edsl_version=add_edsl_version),
+            "model": self.model.to_dict(add_edsl_version=add_edsl_version),
             "iteration": self.iteration,
             "exceptions": {},
         }
@@ -202,11 +202,11 @@ class Interview:
     def __hash__(self) -> int:
         from edsl.utilities.utilities import dict_hash
 
-        return dict_hash(self._to_dict(include_exceptions=False))
+        return dict_hash(self.to_dict(include_exceptions=False, add_edsl_version=False))
 
     def __eq__(self, other: "Interview") -> bool:
         """
-        >>> from edsl.jobs.interviews.Interview import Interview; i = Interview.example(); d = i._to_dict(); i2 = Interview.from_dict(d); i == i2
+        >>> from edsl.jobs.interviews.Interview import Interview; i = Interview.example(); d = i.to_dict(); i2 = Interview.from_dict(d); i == i2
         True
         """
         return hash(self) == hash(other)

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -2,13 +2,14 @@ from __future__ import annotations
 import time
 import asyncio
 import threading
-from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator
-from contextlib import contextmanager
+import warnings
+from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator, Type
+from uuid import UUID
 from collections import UserList
 
 from edsl.results.Results import Results
 from edsl.jobs.interviews.Interview import Interview
-from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
+from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus, JobsRunnerStatusBase
 
 from edsl.jobs.tasks.TaskHistory import TaskHistory
 from edsl.jobs.buckets.BucketCollection import BucketCollection
@@ -36,11 +37,61 @@ class JobsRunnerAsyncio:
     The Jobs object is a collection of interviews that are to be run.
     """
 
+    MAX_CONCURRENT_DEFAULT = 500
+
     def __init__(self, jobs: "Jobs"):
         self.jobs = jobs
         self.interviews: List["Interview"] = jobs.interviews()
         self.bucket_collection: "BucketCollection" = jobs.bucket_collection
         self.total_interviews: List["Interview"] = []
+        self._initialized = threading.Event()
+
+        from edsl.config import CONFIG
+
+        self.MAX_CONCURRENT = int(CONFIG.get("EDSL_MAX_CONCURRENT_TASKS"))
+        # print(f"MAX_CONCURRENT: {self.MAX_CONCURRENT}")
+
+    # async def run_async_generator(
+    #     self,
+    #     cache: Cache,
+    #     n: int = 1,
+    #     stop_on_exception: bool = False,
+    #     sidecar_model: Optional[LanguageModel] = None,
+    #     total_interviews: Optional[List["Interview"]] = None,
+    #     raise_validation_errors: bool = False,
+    # ) -> AsyncGenerator["Result", None]:
+    #     """Creates the tasks, runs them asynchronously, and returns the results as a Results object.
+
+    #     Completed tasks are yielded as they are completed.
+
+    #     :param n: how many times to run each interview
+    #     :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
+    #         self.total_interviews = total_interviews
+    #     else:
+    #         self.total_interviews = list(
+    #             self._populate_total_interviews(n=n)
+    #         )  # Populate self.total_interviews before creating tasks
+    #     self._initialized.set()  # Signal that we're ready
+
+    #     for interview in self.total_interviews:
+    #         interviewing_task = self._build_interview_task(
+    #             interview=interview,
+    #             stop_on_exception=stop_on_exception,
+    #             sidecar_model=sidecar_model,
+    #             raise_validation_errors=raise_validation_errors,
+    #         )
+    #         tasks.append(asyncio.create_task(interviewing_task))
+
+    #     for task in asyncio.as_completed(tasks):
+    #         result = await task
+    #         self.jobs_runner_status.add_completed_interview(result)
+    #         yield result
 
     async def run_async_generator(
         self,
@@ -51,9 +102,10 @@ class JobsRunnerAsyncio:
         total_interviews: Optional[List["Interview"]] = None,
         raise_validation_errors: bool = False,
     ) -> AsyncGenerator["Result", None]:
-        """Creates the tasks, runs them asynchronously, and returns the results as a Results object.
+        """Creates and processes tasks asynchronously, yielding results as they complete.
 
-        Completed tasks are yielded as they are completed.
+        Tasks are created and processed in a streaming fashion rather than building the full list upfront.
+        Results are yielded as soon as they are available.
 
         :param n: how many times to run each interview
         :param stop_on_exception: Whether to stop the interview if an exception is raised
@@ -61,27 +113,70 @@ class JobsRunnerAsyncio:
         :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
+        # Initialize interviews iterator
+        if total_interviews:
+            interviews_iter = iter(total_interviews)
             self.total_interviews = total_interviews
         else:
-            self.total_interviews = list(
-                self._populate_total_interviews(n=n)
-            )  # Populate self.total_interviews before creating tasks
+            interviews_iter = self._populate_total_interviews(n=n)
+            self.total_interviews = list(interviews_iter)
+            interviews_iter = iter(self.total_interviews)  # Create fresh iterator
 
-        for interview in self.total_interviews:
-            interviewing_task = self._build_interview_task(
-                interview=interview,
-                stop_on_exception=stop_on_exception,
-                sidecar_model=sidecar_model,
-                raise_validation_errors=raise_validation_errors,
-            )
-            tasks.append(asyncio.create_task(interviewing_task))
+        self._initialized.set()  # Signal that we're ready
 
-        for task in asyncio.as_completed(tasks):
-            result = await task
-            self.jobs_runner_status.add_completed_interview(result)
-            yield result
+        # Keep track of active tasks
+        active_tasks = set()
+
+        try:
+            while True:
+                # Add new tasks if we're below max_concurrent and there are more interviews
+                while len(active_tasks) < self.MAX_CONCURRENT:
+                    try:
+                        interview = next(interviews_iter)
+                        task = asyncio.create_task(
+                            self._build_interview_task(
+                                interview=interview,
+                                stop_on_exception=stop_on_exception,
+                                sidecar_model=sidecar_model,
+                                raise_validation_errors=raise_validation_errors,
+                            )
+                        )
+                        active_tasks.add(task)
+                        # Add callback to remove task from set when done
+                        task.add_done_callback(active_tasks.discard)
+                    except StopIteration:
+                        break
+
+                if not active_tasks:
+                    break
+
+                # Wait for next completed task
+                done, _ = await asyncio.wait(
+                    active_tasks, return_when=asyncio.FIRST_COMPLETED
+                )
+
+                # Process completed tasks
+                for task in done:
+                    try:
+                        result = await task
+                        self.jobs_runner_status.add_completed_interview(result)
+                        yield result
+                    except Exception as e:
+                        if stop_on_exception:
+                            # Cancel remaining tasks
+                            for t in active_tasks:
+                                if not t.done():
+                                    t.cancel()
+                            raise
+                        else:
+                            # Log error and continue
+                            # logger.error(f"Task failed with error: {e}")
+                            continue
+        finally:
+            # Ensure we cancel any remaining tasks if we exit early
+            for task in active_tasks:
+                if not task.done():
+                    task.cancel()
 
     def _populate_total_interviews(
         self, n: int = 1
@@ -242,11 +337,25 @@ class JobsRunnerAsyncio:
             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
+            import sys
+
+            print(msg, file=sys.stderr)
+            from edsl.config import CONFIG
+
+            if CONFIG.get("EDSL_OPEN_EXCEPTION_REPORT_URL") == "True":
+                open_in_browser = True
+            elif CONFIG.get("EDSL_OPEN_EXCEPTION_REPORT_URL") == "False":
+                open_in_browser = False
+            else:
+                raise Exception(
+                    "EDSL_OPEN_EXCEPTION_REPORT_URL", "must be either True or False"
+                )
+
+            # print("open_in_browser", open_in_browser)
+
             filepath = results.task_history.html(
                 cta="Open report to see details.",
-                open_in_browser=True,
+                open_in_browser=open_in_browser,
                 return_link=True,
             )
 
@@ -275,6 +384,8 @@ class JobsRunnerAsyncio:
         stop_on_exception: bool = False,
         progress_bar: bool = False,
         sidecar_model: Optional[LanguageModel] = None,
+        jobs_runner_status: Optional[Type[JobsRunnerStatusBase]] = None,
+        job_uuid: Optional[UUID] = None,
         print_exceptions: bool = True,
         raise_validation_errors: bool = False,
     ) -> "Coroutine":
@@ -286,7 +397,19 @@ class JobsRunnerAsyncio:
         self.cache = cache
         self.sidecar_model = sidecar_model
 
-        self.jobs_runner_status = JobsRunnerStatus(self, n=n)
+        from edsl.coop import Coop
+
+        coop = Coop()
+        endpoint_url = coop.get_progress_bar_url()
+
+        if jobs_runner_status is not None:
+            self.jobs_runner_status = jobs_runner_status(
+                self, n=n, endpoint_url=endpoint_url, job_uuid=job_uuid
+            )
+        else:
+            self.jobs_runner_status = JobsRunnerStatus(
+                self, n=n, endpoint_url=endpoint_url, job_uuid=job_uuid
+            )
 
         stop_event = threading.Event()
 
@@ -306,11 +429,16 @@ class JobsRunnerAsyncio:
             """Runs the progress bar in a separate thread."""
             self.jobs_runner_status.update_progress(stop_event)
 
-        if progress_bar:
+        if progress_bar and self.jobs_runner_status.has_ep_api_key():
+            self.jobs_runner_status.setup()
             progress_thread = threading.Thread(
                 target=run_progress_bar, args=(stop_event,)
             )
             progress_thread.start()
+        elif progress_bar:
+            warnings.warn(
+                "You need an Expected Parrot API key to view job progress bars."
+            )
 
         exception_to_raise = None
         try:
@@ -325,7 +453,7 @@ class JobsRunnerAsyncio:
             stop_event.set()
         finally:
             stop_event.set()
-            if progress_bar:
+            if progress_bar and self.jobs_runner_status.has_ep_api_key():
                 # self.jobs_runner_status.stop_event.set()
                 if progress_thread:
                     progress_thread.join()