edsl 0.1.39.dev2__py3-none-any.whl → 0.1.39.dev4__py3-none-any.whl

edsl/jobs/JobsPrompts.py

@@ -51,7 +51,7 @@ class JobsPrompts:
         for interview_index, interview in enumerate(interviews):
             invigilators = [
                 FetchInvigilator(interview)(question)
-                for question in self.survey.questions
+                for question in interview.survey.questions
             ]
             for _, invigilator in enumerate(invigilators):
                 prompts = invigilator.get_prompts()

edsl/jobs/JobsRemoteInferenceHandler.py

@@ -1,4 +1,6 @@
-from typing import Optional, Union, Literal, TYPE_CHECKING, NewType
+from typing import Optional, Union, Literal, TYPE_CHECKING, NewType, Callable, Any
+
+from dataclasses import dataclass
 
 
 Seconds = NewType("Seconds", float)
@@ -16,26 +18,52 @@ from edsl.coop.coop import RemoteInferenceResponse, RemoteInferenceCreationInfo
 
 from edsl.jobs.jobs_status_enums import JobsStatus
 from edsl.coop.utils import VisibilityType
+from edsl.jobs.JobsRemoteInferenceLogger import JobLogger
+
+
+class RemoteJobConstants:
+    """Constants for remote job handling."""
+
+    REMOTE_JOB_POLL_INTERVAL = 1
+    REMOTE_JOB_VERBOSE = False
+    DISCORD_URL = "https://discord.com/invite/mxAYkjfy9m"
+
+
+@dataclass
+class RemoteJobInfo:
+    creation_data: RemoteInferenceCreationInfo
+    job_uuid: JobUUID
+    logger: JobLogger
 
 
 class JobsRemoteInferenceHandler:
-    def __init__(self, jobs: "Jobs", verbose: bool = False, poll_interval: Seconds = 1):
-        """ """
+    def __init__(
+        self,
+        jobs: "Jobs",
+        verbose: bool = RemoteJobConstants.REMOTE_JOB_VERBOSE,
+        poll_interval: Seconds = RemoteJobConstants.REMOTE_JOB_POLL_INTERVAL,
+    ):
+        """Handles the creation and running of a remote inference job."""
         self.jobs = jobs
         self.verbose = verbose
         self.poll_interval = poll_interval
 
-        self._remote_job_creation_data: Union[None, RemoteInferenceCreationInfo] = None
-        self._job_uuid: Union[None, JobUUID] = None  # Will be set when job is created
-        self.logger: Union[None, JobLogger] = None  # Will be initialized when needed
+        from edsl.config import CONFIG
 
-    @property
-    def remote_job_creation_data(self) -> RemoteInferenceCreationInfo:
-        return self._remote_job_creation_data
+        self.expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
+        self.remote_inference_url = f"{self.expected_parrot_url}/home/remote-inference"
 
-    @property
-    def job_uuid(self) -> JobUUID:
-        return self._job_uuid
+    def _create_logger(self) -> JobLogger:
+        from edsl.utilities.is_notebook import is_notebook
+        from edsl.jobs.JobsRemoteInferenceLogger import (
+            JupyterJobLogger,
+            StdOutJobLogger,
+        )
+        from edsl.jobs.loggers.HTMLTableJobLogger import HTMLTableJobLogger
+
+        if is_notebook():
+            return HTMLTableJobLogger(verbose=self.verbose)
+        return StdOutJobLogger(verbose=self.verbose)
 
     def use_remote_inference(self, disable_remote_inference: bool) -> bool:
         import requests
@@ -60,23 +88,15 @@ class JobsRemoteInferenceHandler:
         iterations: int = 1,
         remote_inference_description: Optional[str] = None,
         remote_inference_results_visibility: Optional[VisibilityType] = "unlisted",
-    ) -> None:
+    ) -> RemoteJobInfo:
+
         from edsl.config import CONFIG
         from edsl.coop.coop import Coop
 
-        # Initialize logger
-        from edsl.utilities.is_notebook import is_notebook
-        from edsl.jobs.JobsRemoteInferenceLogger import JupyterJobLogger
-        from edsl.jobs.JobsRemoteInferenceLogger import StdOutJobLogger
-        from edsl.jobs.loggers.HTMLTableJobLogger import HTMLTableJobLogger
-
-        if is_notebook():
-            self.logger = HTMLTableJobLogger(verbose=self.verbose)
-        else:
-            self.logger = StdOutJobLogger(verbose=self.verbose)
+        logger = self._create_logger()
 
         coop = Coop()
-        self.logger.update(
+        logger.update(
             "Remote inference activated. Sending job to server...",
             status=JobsStatus.QUEUED,
         )
@@ -87,33 +107,34 @@ class JobsRemoteInferenceHandler:
             iterations=iterations,
             initial_results_visibility=remote_inference_results_visibility,
         )
-        self.logger.update(
+        logger.update(
             "Your survey is running at the Expected Parrot server...",
             status=JobsStatus.RUNNING,
         )
-
         job_uuid = remote_job_creation_data.get("uuid")
-        self.logger.update(
+        logger.update(
             message=f"Job sent to server. (Job uuid={job_uuid}).",
             status=JobsStatus.RUNNING,
         )
-        self.logger.add_info("job_uuid", job_uuid)
+        logger.add_info("job_uuid", job_uuid)
 
-        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
-        remote_inference_url = f"{expected_parrot_url}/home/remote-inference"
-
-        self.logger.update(
-            f"Job details are available at your Coop account {remote_inference_url}{remote_inference_url}",
+        logger.update(
+            f"Job details are available at your Coop account {self.remote_inference_url}",
             status=JobsStatus.RUNNING,
         )
-        progress_bar_url = f"{expected_parrot_url}/home/remote-job-progress/{job_uuid}"
-        self.logger.add_info("progress_bar_url", progress_bar_url)
-        self.logger.update(
+        progress_bar_url = (
+            f"{self.expected_parrot_url}/home/remote-job-progress/{job_uuid}"
+        )
+        logger.add_info("progress_bar_url", progress_bar_url)
+        logger.update(
             f"View job progress here: {progress_bar_url}", status=JobsStatus.RUNNING
         )
 
-        self._remote_job_creation_data = remote_job_creation_data
-        self._job_uuid = job_uuid
+        return RemoteJobInfo(
+            creation_data=remote_job_creation_data,
+            job_uuid=job_uuid,
+            logger=logger,
+        )
 
     @staticmethod
     def check_status(
@@ -124,126 +145,127 @@ class JobsRemoteInferenceHandler:
         coop = Coop()
         return coop.remote_inference_get(job_uuid)
 
-    def poll_remote_inference_job(self) -> Union[None, "Results"]:
-        return self._poll_remote_inference_job(
-            self.remote_job_creation_data, verbose=self.verbose
+    def _construct_remote_job_fetcher(
+        self, testing_simulated_response: Optional[Any] = None
+    ) -> Callable:
+        if testing_simulated_response is not None:
+            return lambda job_uuid: testing_simulated_response
+        else:
+            from edsl.coop.coop import Coop
+
+            coop = Coop()
+            return coop.remote_inference_get
+
+    def _construct_object_fetcher(
+        self, testing_simulated_response: Optional[Any] = None
+    ) -> Callable:
+        "Constructs a function to fetch the results object from Coop."
+        if testing_simulated_response is not None:
+            return lambda results_uuid, expected_object_type: Results.example()
+        else:
+            from edsl.coop.coop import Coop
+
+            coop = Coop()
+            return coop.get
+
+    def _handle_cancelled_job(self, job_info: RemoteJobInfo) -> None:
+        "Handles a cancelled job by logging the cancellation and updating the job status."
+
+        job_info.logger.update(
+            message="Job cancelled by the user.", status=JobsStatus.CANCELLED
+        )
+        job_info.logger.update(
+            f"See {self.expected_parrot_url}/home/remote-inference for more details.",
+            status=JobsStatus.CANCELLED,
         )
 
-    def _poll_remote_inference_job(
-        self,
-        remote_job_creation_data: RemoteInferenceCreationInfo,
-        verbose: bool = False,
-        poll_interval: Optional[Seconds] = None,
-        testing_simulated_response=None,
-    ) -> Union[None, "Results"]:
+    def _handle_failed_job(
+        self, job_info: RemoteJobInfo, remote_job_data: RemoteInferenceResponse
+    ) -> None:
+        "Handles a failed job by logging the error and updating the job status."
+        latest_error_report_url = remote_job_data.get("latest_error_report_url")
+        if latest_error_report_url:
+            job_info.logger.add_info("error_report_url", latest_error_report_url)
+
+        job_info.logger.update("Job failed.", status=JobsStatus.FAILED)
+        job_info.logger.update(
+            f"See {self.expected_parrot_url}/home/remote-inference for more details.",
+            status=JobsStatus.FAILED,
+        )
+        job_info.logger.update(
+            f"Need support? Visit Discord: {RemoteJobConstants.DISCORD_URL}",
+            status=JobsStatus.FAILED,
+        )
+
+    def _sleep_for_a_bit(self, job_info: RemoteJobInfo, status: str) -> None:
         import time
         from datetime import datetime
-        from edsl.config import CONFIG
-        from edsl.results.Results import Results
 
-        if poll_interval is None:
-            poll_interval = self.poll_interval
+        time_checked = datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")
+        job_info.logger.update(
+            f"Job status: {status} - last update: {time_checked}",
+            status=JobsStatus.RUNNING,
+        )
+        time.sleep(self.poll_interval)
 
-        job_uuid = remote_job_creation_data.get("uuid")
-        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
+    def _fetch_results_and_log(
+        self,
+        job_info: RemoteJobInfo,
+        results_uuid: str,
+        remote_job_data: RemoteInferenceResponse,
+        object_fetcher: Callable,
+    ) -> "Results":
+        "Fetches the results object and logs the results URL."
+        job_info.logger.add_info("results_uuid", results_uuid)
+        results = object_fetcher(results_uuid, expected_object_type="results")
+        results_url = remote_job_data.get("results_url")
+        job_info.logger.update(
+            f"Job completed and Results stored on Coop: {results_url}",
+            status=JobsStatus.COMPLETED,
+        )
+        results.job_uuid = job_info.job_uuid
+        results.results_uuid = results_uuid
+        return results
 
-        if testing_simulated_response is not None:
-            remote_job_data_fetcher = lambda job_uuid: testing_simulated_response
-            object_fetcher = (
-                lambda results_uuid, expected_object_type: Results.example()
-            )
-        else:
-            from edsl.coop.coop import Coop
+    def poll_remote_inference_job(
+        self,
+        job_info: RemoteJobInfo,
+        testing_simulated_response=None,
+    ) -> Union[None, "Results"]:
+        """Polls a remote inference job for completion and returns the results."""
 
-            coop = Coop()
-            remote_job_data_fetcher = coop.remote_inference_get
-            object_fetcher = coop.get
+        remote_job_data_fetcher = self._construct_remote_job_fetcher(
+            testing_simulated_response
+        )
+        object_fetcher = self._construct_object_fetcher(testing_simulated_response)
 
         job_in_queue = True
         while job_in_queue:
-            remote_job_data: RemoteInferenceResponse = remote_job_data_fetcher(job_uuid)
+            remote_job_data = remote_job_data_fetcher(job_info.job_uuid)
             status = remote_job_data.get("status")
 
             if status == "cancelled":
-                self.logger.update(
-                    messaged="Job cancelled by the user.", status=JobsStatus.CANCELLED
-                )
-                self.logger.update(
-                    f"See {expected_parrot_url}/home/remote-inference for more details.",
-                    status=JobsStatus.CANCELLED,
-                )
+                self._handle_cancelled_job(job_info)
                 return None
 
-            elif status == "failed":
-                latest_error_report_url = remote_job_data.get("latest_error_report_url")
-                if latest_error_report_url:
-                    self.logger.update("Job failed.", status=JobsStatus.FAILED)
-                    self.logger.update(
-                        f"Error report: {latest_error_report_url}", "failed"
-                    )
-                    self.logger.add_info("error_report_url", latest_error_report_url)
-                    self.logger.update(
-                        "Need support? Visit Discord: https://discord.com/invite/mxAYkjfy9m",
-                        status=JobsStatus.FAILED,
-                    )
-                else:
-                    self.logger.update("Job failed.", "failed")
-                    self.logger.update(
-                        f"See {expected_parrot_url}/home/remote-inference for details.",
-                        status=JobsStatus.FAILED,
-                    )
+            elif status == "failed" or status == "completed":
+                if status == "failed":
+                    self._handle_failed_job(job_info, remote_job_data)
 
                 results_uuid = remote_job_data.get("results_uuid")
                 if results_uuid:
-                    self.logger.add_info("results_uuid", results_uuid)
-                    results = object_fetcher(
-                        results_uuid, expected_object_type="results"
+                    results = self._fetch_results_and_log(
+                        job_info=job_info,
+                        results_uuid=results_uuid,
+                        remote_job_data=remote_job_data,
+                        object_fetcher=object_fetcher,
                     )
-                    results.job_uuid = job_uuid
-                    results.results_uuid = results_uuid
                     return results
                 else:
                     return None
 
-            elif status == "completed":
-                results_uuid = remote_job_data.get("results_uuid")
-                self.logger.add_info("results_uuid", results_uuid)
-                results_url = remote_job_data.get("results_url")
-                self.logger.add_info("results_url", results_url)
-                results = object_fetcher(results_uuid, expected_object_type="results")
-                self.logger.update(
-                    f"Job completed and Results stored on Coop: {results_url}",
-                    status=JobsStatus.COMPLETED,
-                )
-                results.job_uuid = job_uuid
-                results.results_uuid = results_uuid
-                return results
-
             else:
-                time_checked = datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")
-                self.logger.update(
-                    f"Job status: {status} - last update: {time_checked}",
-                    status=JobsStatus.RUNNING,
-                )
-                time.sleep(poll_interval)
-
-    def use_remote_inference(self, disable_remote_inference: bool) -> bool:
-        import requests
-
-        if disable_remote_inference:
-            return False
-        if not disable_remote_inference:
-            try:
-                from edsl.coop.coop import Coop
-
-                user_edsl_settings = Coop().edsl_settings
-                return user_edsl_settings.get("remote_inference", False)
-            except requests.ConnectionError:
-                pass
-            except CoopServerResponseError as e:
-                pass
-
-        return False
+                self._sleep_for_a_bit(job_info, status)
 
     async def create_and_poll_remote_job(
         self,
@@ -265,7 +287,7 @@ class JobsRemoteInferenceHandler:
 
         # Create job using existing method
         loop = asyncio.get_event_loop()
-        remote_job_creation_data = await loop.run_in_executor(
+        job_info = await loop.run_in_executor(
             None,
             partial(
                 self.create_remote_inference_job,
@@ -274,10 +296,12 @@ class JobsRemoteInferenceHandler:
                 remote_inference_results_visibility=remote_inference_results_visibility,
             ),
         )
+        if job_info is None:
+            raise ValueError("Remote job creation failed.")
 
-        # Poll using existing method but with async sleep
         return await loop.run_in_executor(
-            None, partial(self.poll_remote_inference_job, remote_job_creation_data)
+            None,
+            partial(self.poll_remote_inference_job, job_info),
         )
 
 

edsl/jobs/async_interview_runner.py

@@ -0,0 +1,138 @@
+from collections.abc import AsyncGenerator
+from typing import List, TypeVar, Generator, Tuple, TYPE_CHECKING
+from dataclasses import dataclass
+import asyncio
+from contextlib import asynccontextmanager
+from edsl.data_transfer_models import EDSLResultObjectInput
+
+from edsl.results.Result import Result
+from edsl.jobs.interviews.Interview import Interview
+
+if TYPE_CHECKING:
+    from edsl.jobs.Jobs import Jobs
+
+
+@dataclass
+class InterviewResult:
+    result: Result
+    interview: Interview
+    order: int
+
+
+from edsl.jobs.data_structures import RunConfig
+
+
+class AsyncInterviewRunner:
+    MAX_CONCURRENT = 5
+
+    def __init__(self, jobs: "Jobs", run_config: RunConfig):
+        self.jobs = jobs
+        self.run_config = run_config
+        self._initialized = asyncio.Event()
+
+    def _expand_interviews(self) -> Generator["Interview", None, None]:
+        """Populates self.total_interviews with n copies of each interview.
+
+        It also has to set the cache for each interview.
+
+        :param n: how many times to run each interview.
+        """
+        for interview in self.jobs.generate_interviews():
+            for iteration in range(self.run_config.parameters.n):
+                if iteration > 0:
+                    yield interview.duplicate(
+                        iteration=iteration, cache=self.run_config.environment.cache
+                    )
+                else:
+                    interview.cache = self.run_config.environment.cache
+                    yield interview
+
+    async def _conduct_interview(
+        self, interview: "Interview"
+    ) -> Tuple["Result", "Interview"]:
+        """Conducts an interview and returns the result object, along with the associated interview.
+
+        We return the interview because it is not populated with exceptions, if any.
+
+        :param interview: the interview to conduct
+        :return: the result of the interview
+
+        'extracted_answers' is a dictionary of the answers to the questions in the interview.
+        This is not the same as the generated_tokens---it can include substantial cleaning and processing / validation.
+        """
+        # the model buckets are used to track usage rates
+        # model_buckets = self.bucket_collection[interview.model]
+        # model_buckets = self.run_config.environment.bucket_collection[interview.model]
+
+        # get the results of the interview e.g., {'how_are_you':"Good" 'how_are_you_generated_tokens': "Good"}
+        extracted_answers: dict[str, str]
+        model_response_objects: List[EDSLResultObjectInput]
+
+        extracted_answers, model_response_objects = (
+            await interview.async_conduct_interview(self.run_config)
+        )
+        result = Result.from_interview(
+            interview=interview,
+            extracted_answers=extracted_answers,
+            model_response_objects=model_response_objects,
+        )
+        return result, interview
+
+    async def run(
+        self,
+    ) -> AsyncGenerator[tuple[Result, Interview], None]:
+        """Creates and processes tasks asynchronously, yielding results as they complete.
+
+        Uses TaskGroup for structured concurrency and automated cleanup.
+        Results are yielded as they become available while maintaining controlled concurrency.
+        """
+        interviews = list(self._expand_interviews())
+        self._initialized.set()
+
+        async def _process_single_interview(
+            interview: Interview, idx: int
+        ) -> InterviewResult:
+            try:
+                result, interview = await self._conduct_interview(interview)
+                self.run_config.environment.jobs_runner_status.add_completed_interview(
+                    result
+                )
+                result.order = idx
+                return InterviewResult(result, interview, idx)
+            except Exception as e:
+                # breakpoint()
+                if self.run_config.parameters.stop_on_exception:
+                    raise
+                # logger.error(f"Task failed with error: {e}")
+                return None
+
+        # Process interviews in chunks
+        for i in range(0, len(interviews), self.MAX_CONCURRENT):
+            chunk = interviews[i : i + self.MAX_CONCURRENT]
+            tasks = [
+                asyncio.create_task(_process_single_interview(interview, idx))
+                for idx, interview in enumerate(chunk, start=i)
+            ]
+
+            try:
+                # Wait for all tasks in the chunk to complete
+                results = await asyncio.gather(
+                    *tasks,
+                    return_exceptions=not self.run_config.parameters.stop_on_exception
+                )
+
+                # Process successful results
+                for result in (r for r in results if r is not None):
+                    yield result.result, result.interview
+
+            except Exception as e:
+                if self.run_config.parameters.stop_on_exception:
+                    raise
+                # logger.error(f"Chunk processing failed with error: {e}")
+                continue
+
+            finally:
+                # Clean up any remaining tasks
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()

edsl/jobs/check_survey_scenario_compatibility.py

@@ -0,0 +1,85 @@
+import warnings
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.surveys.Survey import Survey
+    from edsl.scenarios.ScenarioList import ScenarioList
+
+
+class CheckSurveyScenarioCompatibility:
+
+    def __init__(self, survey: "Survey", scenarios: "ScenarioList"):
+        self.survey = survey
+        self.scenarios = scenarios
+
+    def check(self, strict: bool = False, warn: bool = False) -> None:
+        """Check if the parameters in the survey and scenarios are consistent.
+
+        >>> from edsl.jobs.Jobs import Jobs
+        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
+        >>> from edsl.surveys.Survey import Survey
+        >>> from edsl.scenarios.Scenario import Scenario
+        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
+        >>> j = Jobs(survey = Survey(questions=[q]))
+        >>> cs = CheckSurveyScenarioCompatibility(j.survey, j.scenarios)
+        >>> with warnings.catch_warnings(record=True) as w:
+        ...     cs.check(warn = True)
+        ...     assert len(w) == 1
+        ...     assert issubclass(w[-1].category, UserWarning)
+        ...     assert "The following parameters are in the survey but not in the scenarios" in str(w[-1].message)
+
+        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
+        >>> s = Scenario({'plop': "A", 'poo': "B"})
+        >>> j = Jobs(survey = Survey(questions=[q])).by(s)
+        >>> cs = CheckSurveyScenarioCompatibility(j.survey, j.scenarios)
+        >>> cs.check(strict = True)
+        Traceback (most recent call last):
+        ...
+        ValueError: The following parameters are in the scenarios but not in the survey: {'plop'}
+
+        >>> q = QuestionFreeText(question_text = "Hello", question_name = "ugly_question")
+        >>> s = Scenario({'ugly_question': "B"})
+        >>> from edsl.scenarios.ScenarioList import ScenarioList
+        >>> cs = CheckSurveyScenarioCompatibility(Survey(questions=[q]), ScenarioList([s]))
+        >>> cs.check()
+        Traceback (most recent call last):
+        ...
+        ValueError: The following names are in both the survey question_names and the scenario keys: {'ugly_question'}. This will create issues.
+        """
+        survey_parameters: set = self.survey.parameters
+        scenario_parameters: set = self.scenarios.parameters
+
+        msg0, msg1, msg2 = None, None, None
+
+        # look for key issues
+        if intersection := set(self.scenarios.parameters) & set(
+            self.survey.question_names
+        ):
+            msg0 = f"The following names are in both the survey question_names and the scenario keys: {intersection}. This will create issues."
+
+            raise ValueError(msg0)
+
+        if in_survey_but_not_in_scenarios := survey_parameters - scenario_parameters:
+            msg1 = f"The following parameters are in the survey but not in the scenarios: {in_survey_but_not_in_scenarios}"
+        if in_scenarios_but_not_in_survey := scenario_parameters - survey_parameters:
+            msg2 = f"The following parameters are in the scenarios but not in the survey: {in_scenarios_but_not_in_survey}"
+
+        if msg1 or msg2:
+            message = "\n".join(filter(None, [msg1, msg2]))
+            if strict:
+                raise ValueError(message)
+            else:
+                if warn:
+                    warnings.warn(message)
+
+        if self.scenarios.has_jinja_braces:
+            warnings.warn(
+                "The scenarios have Jinja braces ({{ and }}). Converting to '<<' and '>>'. If you want a different conversion, use the convert_jinja_braces method first to modify the scenario."
+            )
+            self.scenarios = self.scenarios._convert_jinja_braces()
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()