edsl 0.1.39__py3-none-any.whl → 0.1.39.dev1__py3-none-any.whl

edsl/jobs/JobsChecks.py

@@ -1,21 +1,16 @@
 import os
-from edsl.exceptions.general import MissingAPIKeyError
+from edsl.exceptions import MissingAPIKeyError
 
 
 class JobsChecks:
     def __init__(self, jobs):
-        """Checks a Jobs object for missing API keys and other requirements."""
+        """ """
         self.jobs = jobs
 
     def check_api_keys(self) -> None:
-        from edsl.language_models.model import Model
+        from edsl import Model
 
-        if len(self.jobs.models) == 0:
-            models = [Model()]
-        else:
-            models = self.jobs.models
-
-        for model in models:  # + [Model()]:
+        for model in self.jobs.models + [Model()]:
             if not model.has_valid_api_key():
                 raise MissingAPIKeyError(
                     model_name=str(model.model),
@@ -28,7 +23,7 @@ class JobsChecks:
         """
         missing_api_keys = set()
 
-        from edsl.language_models.model import Model
+        from edsl import Model
         from edsl.enums import service_to_api_keyname
 
         for model in self.jobs.models + [Model()]:
@@ -100,33 +95,16 @@ class JobsChecks:
             return True
 
     def needs_key_process(self):
-        """
-        A User needs the key process when:
-        1. They don't have all the model keys
-        2. They don't have the EP API
-        3. They need external LLMs to run the job
-        """
         return (
             not self.user_has_all_model_keys()
             and not self.user_has_ep_api_key()
             and self.needs_external_llms()
         )
 
-    def status(self) -> dict:
-        """
-        Returns a dictionary with the status of the job checks.
-        """
-        return {
-            "user_has_ep_api_key": self.user_has_ep_api_key(),
-            "user_has_all_model_keys": self.user_has_all_model_keys(),
-            "needs_external_llms": self.needs_external_llms(),
-            "needs_key_process": self.needs_key_process(),
-        }
-
     def key_process(self):
         import secrets
         from dotenv import load_dotenv
-        from edsl.config import CONFIG
+        from edsl import CONFIG
         from edsl.coop.coop import Coop
         from edsl.utilities.utilities import write_api_key_to_env
 
@@ -141,12 +119,10 @@ class JobsChecks:
             "\nYou can either add the missing keys to your .env file, or use remote inference."
         )
         print("Remote inference allows you to run jobs on our server.")
+        print("\n🚀 To use remote inference, sign up at the following link:")
 
         coop = Coop()
-        coop._display_login_url(
-            edsl_auth_token=edsl_auth_token,
-            link_description="\n🚀 To use remote inference, sign up at the following link:",
-        )
+        coop._display_login_url(edsl_auth_token=edsl_auth_token)
 
         print(
             "\nOnce you log in, we will automatically retrieve your Expected Parrot API key and continue your job remotely."
@@ -158,9 +134,8 @@ class JobsChecks:
             print("\nTimed out waiting for login. Please try again.")
             return
 
-        path_to_env = write_api_key_to_env(api_key)
-        print("\n✨ API key retrieved and written to .env file at the following path:")
-        print(f"    {path_to_env}")
+        write_api_key_to_env(api_key)
+        print("✨ API key retrieved and written to .env file.\n")
 
         # Retrieve API key so we can continue running the job
         load_dotenv()

edsl/jobs/JobsPrompts.py

@@ -11,8 +11,6 @@ if TYPE_CHECKING:
     # from edsl.scenarios.ScenarioList import ScenarioList
     # from edsl.surveys.Survey import Survey
 
-from edsl.jobs.FetchInvigilator import FetchInvigilator
-
 
 class JobsPrompts:
     def __init__(self, jobs: "Jobs"):
@@ -25,7 +23,7 @@ class JobsPrompts:
     @property
     def price_lookup(self):
         if self._price_lookup is None:
-            from edsl.coop.coop import Coop
+            from edsl import Coop
 
             c = Coop()
             self._price_lookup = c.fetch_prices()
@@ -50,8 +48,8 @@ class JobsPrompts:
 
         for interview_index, interview in enumerate(interviews):
             invigilators = [
-                FetchInvigilator(interview)(question)
-                for question in interview.survey.questions
+                interview._get_invigilator(question)
+                for question in self.survey.questions
             ]
             for _, invigilator in enumerate(invigilators):
                 prompts = invigilator.get_prompts()
@@ -186,7 +184,7 @@ class JobsPrompts:
         data = []
         for interview in interviews:
             invigilators = [
-                FetchInvigilator(interview)(question)
+                interview._get_invigilator(question)
                 for question in self.survey.questions
             ]
             for invigilator in invigilators:

edsl/jobs/JobsRemoteInferenceHandler.py

@@ -1,78 +1,47 @@
-from typing import Optional, Union, Literal, TYPE_CHECKING, NewType, Callable, Any
-
-from dataclasses import dataclass
-
-
-Seconds = NewType("Seconds", float)
-JobUUID = NewType("JobUUID", str)
-
+from typing import Optional, Union, Literal
+import requests
+import sys
 from edsl.exceptions.coop import CoopServerResponseError
 
-if TYPE_CHECKING:
-    from edsl.results.Results import Results
-    from edsl.jobs.Jobs import Jobs
-    from edsl.coop.coop import RemoteInferenceResponse, RemoteInferenceCreationInfo
-    from edsl.jobs.JobsRemoteInferenceLogger import JobLogger
-
-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
+# from edsl.enums import VisibilityType
+from edsl.results import Results
 
 
 class JobsRemoteInferenceHandler:
-    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."""
+    def __init__(self, jobs, verbose=False, poll_interval=3):
+        """
+        >>> from edsl.jobs import Jobs
+        >>> jh = JobsRemoteInferenceHandler(Jobs.example(), verbose=True)
+        >>> jh.use_remote_inference(True)
+        False
+        >>> jh._poll_remote_inference_job({'uuid':1234}, testing_simulated_response={"status": "failed"}) # doctest: +NORMALIZE_WHITESPACE
+        Job failed.
+        ...
+        >>> jh._poll_remote_inference_job({'uuid':1234}, testing_simulated_response={"status": "completed"}) # doctest: +NORMALIZE_WHITESPACE
+        Job completed and Results stored on Coop: None.
+        Results(...)
+        """
         self.jobs = jobs
         self.verbose = verbose
         self.poll_interval = poll_interval
 
-        from edsl.config import CONFIG
-
-        self.expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
-        self.remote_inference_url = f"{self.expected_parrot_url}/home/remote-inference"
+        self._remote_job_creation_data = None
+        self._job_uuid = None
 
-    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
+    @property
+    def remote_job_creation_data(self):
+        return self._remote_job_creation_data
 
-        if is_notebook():
-            return HTMLTableJobLogger(verbose=self.verbose)
-        return StdOutJobLogger(verbose=self.verbose)
+    @property
+    def job_uuid(self):
+        return self._job_uuid
 
     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
+                from edsl import Coop
 
                 user_edsl_settings = Coop().edsl_settings
                 return user_edsl_settings.get("remote_inference", False)
@@ -87,19 +56,16 @@ class JobsRemoteInferenceHandler:
         self,
         iterations: int = 1,
         remote_inference_description: Optional[str] = None,
-        remote_inference_results_visibility: Optional[VisibilityType] = "unlisted",
-    ) -> RemoteJobInfo:
-
+        remote_inference_results_visibility: Optional["VisibilityType"] = "unlisted",
+        verbose=False,
+    ):
+        """ """
         from edsl.config import CONFIG
         from edsl.coop.coop import Coop
-
-        logger = self._create_logger()
+        from rich import print as rich_print
 
         coop = Coop()
-        logger.update(
-            "Remote inference activated. Sending job to server...",
-            status=JobsStatus.QUEUED,
-        )
+        print("Remote inference activated. Sending job to server...")
         remote_job_creation_data = coop.remote_inference_create(
             self.jobs,
             description=remote_inference_description,
@@ -107,172 +73,136 @@ class JobsRemoteInferenceHandler:
             iterations=iterations,
             initial_results_visibility=remote_inference_results_visibility,
         )
-        logger.update(
-            "Your survey is running at the Expected Parrot server...",
-            status=JobsStatus.RUNNING,
-        )
         job_uuid = remote_job_creation_data.get("uuid")
-        logger.update(
-            message=f"Job sent to server. (Job uuid={job_uuid}).",
-            status=JobsStatus.RUNNING,
-        )
-        logger.add_info("job_uuid", job_uuid)
+        print(f"Job sent to server. (Job uuid={job_uuid}).")
 
-        logger.update(
-            f"Job details are available at your Coop account {self.remote_inference_url}",
-            status=JobsStatus.RUNNING,
-        )
-        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
-        )
+        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
+        progress_bar_url = f"{expected_parrot_url}/home/remote-job-progress/{job_uuid}"
 
-        return RemoteJobInfo(
-            creation_data=remote_job_creation_data,
-            job_uuid=job_uuid,
-            logger=logger,
+        rich_print(
+            f"View job progress here: [#38bdf8][link={progress_bar_url}]{progress_bar_url}[/link][/#38bdf8]"
         )
 
+        self._remote_job_creation_data = remote_job_creation_data
+        self._job_uuid = job_uuid
+        # return remote_job_creation_data
+
     @staticmethod
-    def check_status(
-        job_uuid: JobUUID,
-    ) -> RemoteInferenceResponse:
+    def check_status(job_uuid):
         from edsl.coop.coop import Coop
 
         coop = Coop()
         return coop.remote_inference_get(job_uuid)
 
-    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 _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 poll_remote_inference_job(self):
+        return self._poll_remote_inference_job(
+            self.remote_job_creation_data, verbose=self.verbose
         )
 
-    def _sleep_for_a_bit(self, job_info: RemoteJobInfo, status: str) -> None:
+    def _poll_remote_inference_job(
+        self,
+        remote_job_creation_data: dict,
+        verbose=False,
+        poll_interval: Optional[float] = None,
+        testing_simulated_response: Optional[dict] = None,
+    ) -> Union[Results, None]:
         import time
         from datetime import datetime
+        from edsl.config import CONFIG
+        from edsl.coop.coop import Coop
 
-        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)
+        if poll_interval is None:
+            poll_interval = self.poll_interval
 
-    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
+        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
 
-    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."""
+        job_uuid = remote_job_creation_data.get("uuid")
+        coop = Coop()
 
-        remote_job_data_fetcher = self._construct_remote_job_fetcher(
-            testing_simulated_response
-        )
-        object_fetcher = self._construct_object_fetcher(testing_simulated_response)
+        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:
+            remote_job_data_fetcher = coop.remote_inference_get
+            object_fetcher = coop.get
 
         job_in_queue = True
         while job_in_queue:
-            remote_job_data = remote_job_data_fetcher(job_info.job_uuid)
+            remote_job_data = remote_job_data_fetcher(job_uuid)
             status = remote_job_data.get("status")
-
             if status == "cancelled":
-                self._handle_cancelled_job(job_info)
+                print("\r" + " " * 80 + "\r", end="")
+                print("Job cancelled by the user.")
+                print(
+                    f"See {expected_parrot_url}/home/remote-inference for more details."
+                )
                 return None
-
-            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:
-                    results = self._fetch_results_and_log(
-                        job_info=job_info,
-                        results_uuid=results_uuid,
-                        remote_job_data=remote_job_data,
-                        object_fetcher=object_fetcher,
+            elif status == "failed":
+                print("\r" + " " * 80 + "\r", end="")
+                # write to stderr
+                latest_error_report_url = remote_job_data.get("latest_error_report_url")
+                if latest_error_report_url:
+                    print("Job failed.")
+                    print(
+                        f"Your job generated exceptions. Details on these exceptions can be found in the following report: {latest_error_report_url}"
+                    )
+                    print(
+                        f"Need support? Post a message at the Expected Parrot Discord channel (https://discord.com/invite/mxAYkjfy9m) or send an email to info@expectedparrot.com."
                     )
-                    return results
                 else:
-                    return None
-
+                    print("Job failed.")
+                    print(
+                        f"See {expected_parrot_url}/home/remote-inference for more details."
+                    )
+                return None
+            elif status == "completed":
+                results_uuid = remote_job_data.get("results_uuid")
+                results_url = remote_job_data.get("results_url")
+                results = object_fetcher(results_uuid, expected_object_type="results")
+                print("\r" + " " * 80 + "\r", end="")
+                print(f"Job completed and Results stored on Coop: {results_url}.")
+                return results
             else:
-                self._sleep_for_a_bit(job_info, status)
+                duration = poll_interval
+                time_checked = datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")
+                frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+                start_time = time.time()
+                i = 0
+                while time.time() - start_time < duration:
+                    print(
+                        f"\r{frames[i % len(frames)]} Job status: {status} - last update: {time_checked}",
+                        end="",
+                        flush=True,
+                    )
+                    time.sleep(0.1)
+                    i += 1
+
+    def use_remote_inference(self, disable_remote_inference: bool) -> bool:
+        if disable_remote_inference:
+            return False
+        if not disable_remote_inference:
+            try:
+                from edsl 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
 
     async def create_and_poll_remote_job(
         self,
         iterations: int = 1,
         remote_inference_description: Optional[str] = None,
-        remote_inference_results_visibility: Optional[VisibilityType] = "unlisted",
-    ) -> Union["Results", None]:
+        remote_inference_results_visibility: Optional[
+            Literal["private", "public", "unlisted"]
+        ] = "unlisted",
+    ) -> Union[Results, None]:
         """
         Creates and polls a remote inference job asynchronously.
         Reuses existing synchronous methods but runs them in an async context.
@@ -287,7 +217,7 @@ class JobsRemoteInferenceHandler:
 
         # Create job using existing method
         loop = asyncio.get_event_loop()
-        job_info = await loop.run_in_executor(
+        remote_job_creation_data = await loop.run_in_executor(
             None,
             partial(
                 self.create_remote_inference_job,
@@ -296,12 +226,10 @@ 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, job_info),
+            None, partial(self.poll_remote_inference_job, remote_job_creation_data)
         )
 
 

edsl/jobs/buckets/BucketCollection.py

@@ -1,15 +1,8 @@
-from typing import Optional
 from collections import UserDict
 from edsl.jobs.buckets.TokenBucket import TokenBucket
 from edsl.jobs.buckets.ModelBuckets import ModelBuckets
 
-# from functools import wraps
-from threading import RLock
 
-from edsl.jobs.decorators import synchronized_class
-
-
-@synchronized_class
 class BucketCollection(UserDict):
     """A Jobs object will have a whole collection of model buckets, as multiple models could be used.
 
@@ -17,43 +10,11 @@ class BucketCollection(UserDict):
     Models themselves are hashable, so this works.
     """
 
-    def __init__(self, infinity_buckets: bool = False):
-        """Create a new BucketCollection.
-        An infinity bucket is a bucket that never runs out of tokens or requests.
-        """
+    def __init__(self, infinity_buckets=False):
         super().__init__()
         self.infinity_buckets = infinity_buckets
         self.models_to_services = {}
         self.services_to_buckets = {}
-        self._lock = RLock()
-
-        from edsl.config import CONFIG
-        import os
-
-        url = os.environ.get("EDSL_REMOTE_TOKEN_BUCKET_URL", None)
-
-        if url == "None" or url is None:
-            self.remote_url = None
-            # print(f"Using remote token bucket URL: {url}")
-        else:
-            self.remote_url = url
-
-    @classmethod
-    def from_models(
-        cls, models_list: list, infinity_buckets: bool = False
-    ) -> "BucketCollection":
-        """Create a BucketCollection from a list of models."""
-        bucket_collection = cls(infinity_buckets=infinity_buckets)
-        for model in models_list:
-            bucket_collection.add_model(model)
-        return bucket_collection
-
-    def get_tokens(
-        self, model: "LanguageModel", bucket_type: str, num_tokens: int
-    ) -> int:
-        """Get the number of tokens remaining in the bucket."""
-        relevant_bucket = getattr(self[model], bucket_type)
-        return relevant_bucket.get_tokens(num_tokens)
 
     def __repr__(self):
         return f"BucketCollection({self.data})"
@@ -65,8 +26,8 @@ class BucketCollection(UserDict):
 
         # compute the TPS and RPS from the model
         if not self.infinity_buckets:
-            TPS = model.tpm / 60.0
-            RPS = model.rpm / 60.0
+            TPS = model.TPM / 60.0
+            RPS = model.RPM / 60.0
         else:
             TPS = float("inf")
             RPS = float("inf")
@@ -79,14 +40,12 @@ class BucketCollection(UserDict):
                     bucket_type="requests",
                     capacity=RPS,
                     refill_rate=RPS,
-                    remote_url=self.remote_url,
                 )
                 tokens_bucket = TokenBucket(
                     bucket_name=service,
                     bucket_type="tokens",
                     capacity=TPS,
                     refill_rate=TPS,
-                    remote_url=self.remote_url,
                 )
                 self.services_to_buckets[service] = ModelBuckets(
                     requests_bucket, tokens_bucket