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

edsl/jobs/JobsRemoteInferenceHandler.py

@@ -1,239 +1,311 @@
-from typing import Optional, Union, Literal
-import requests
-import sys
-from edsl.exceptions.coop import CoopServerResponseError
-
-# from edsl.enums import VisibilityType
-from edsl.results import Results
-
-
-class JobsRemoteInferenceHandler:
-    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
-
-        self._remote_job_creation_data = None
-        self._job_uuid = None
-
-    @property
-    def remote_job_creation_data(self):
-        return self._remote_job_creation_data
-
-    @property
-    def job_uuid(self):
-        return self._job_uuid
-
-    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
-
-    def create_remote_inference_job(
-        self,
-        iterations: int = 1,
-        remote_inference_description: Optional[str] = None,
-        remote_inference_results_visibility: Optional["VisibilityType"] = "unlisted",
-        verbose=False,
-    ):
-        """ """
-        from edsl.config import CONFIG
-        from edsl.coop.coop import Coop
-        from rich import print as rich_print
-
-        coop = Coop()
-        print("Remote inference activated. Sending job to server...")
-        remote_job_creation_data = coop.remote_inference_create(
-            self.jobs,
-            description=remote_inference_description,
-            status="queued",
-            iterations=iterations,
-            initial_results_visibility=remote_inference_results_visibility,
-        )
-        job_uuid = remote_job_creation_data.get("uuid")
-        print(f"Job sent to server. (Job uuid={job_uuid}).")
-
-        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
-        progress_bar_url = f"{expected_parrot_url}/home/remote-job-progress/{job_uuid}"
-
-        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):
-        from edsl.coop.coop import Coop
-
-        coop = Coop()
-        return coop.remote_inference_get(job_uuid)
-
-    def poll_remote_inference_job(self):
-        return self._poll_remote_inference_job(
-            self.remote_job_creation_data, verbose=self.verbose
-        )
-
-    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
-
-        if poll_interval is None:
-            poll_interval = self.poll_interval
-
-        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
-
-        job_uuid = remote_job_creation_data.get("uuid")
-        coop = Coop()
-
-        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_uuid)
-            status = remote_job_data.get("status")
-            if status == "cancelled":
-                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":
-                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."
-                    )
-                else:
-                    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:
-                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[
-            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.
-
-        :param iterations: Number of times to run each interview
-        :param remote_inference_description: Optional description for the remote job
-        :param remote_inference_results_visibility: Visibility setting for results
-        :return: Results object if successful, None if job fails or is cancelled
-        """
-        import asyncio
-        from functools import partial
-
-        # Create job using existing method
-        loop = asyncio.get_event_loop()
-        remote_job_creation_data = await loop.run_in_executor(
-            None,
-            partial(
-                self.create_remote_inference_job,
-                iterations=iterations,
-                remote_inference_description=remote_inference_description,
-                remote_inference_results_visibility=remote_inference_results_visibility,
-            ),
-        )
-
-        # 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)
-        )
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod(optionflags=doctest.ELLIPSIS)
+from typing import Optional, Union, Literal, TYPE_CHECKING, NewType, Callable, Any
+
+from dataclasses import dataclass
+
+
+Seconds = NewType("Seconds", float)
+JobUUID = NewType("JobUUID", str)
+
+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
+
+
+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."""
+        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"
+
+    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
+
+        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
+
+    def create_remote_inference_job(
+        self,
+        iterations: int = 1,
+        remote_inference_description: Optional[str] = None,
+        remote_inference_results_visibility: Optional[VisibilityType] = "unlisted",
+    ) -> RemoteJobInfo:
+
+        from edsl.config import CONFIG
+        from edsl.coop.coop import Coop
+
+        logger = self._create_logger()
+
+        coop = Coop()
+        logger.update(
+            "Remote inference activated. Sending job to server...",
+            status=JobsStatus.QUEUED,
+        )
+        remote_job_creation_data = coop.remote_inference_create(
+            self.jobs,
+            description=remote_inference_description,
+            status="queued",
+            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)
+
+        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
+        )
+
+        return RemoteJobInfo(
+            creation_data=remote_job_creation_data,
+            job_uuid=job_uuid,
+            logger=logger,
+        )
+
+    @staticmethod
+    def check_status(
+        job_uuid: JobUUID,
+    ) -> RemoteInferenceResponse:
+        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 _sleep_for_a_bit(self, job_info: RemoteJobInfo, status: str) -> None:
+        import time
+        from datetime import datetime
+
+        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)
+
+    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
+
+    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."""
+
+        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 = remote_job_data_fetcher(job_info.job_uuid)
+            status = remote_job_data.get("status")
+
+            if status == "cancelled":
+                self._handle_cancelled_job(job_info)
+                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,
+                    )
+                    return results
+                else:
+                    return None
+
+            else:
+                self._sleep_for_a_bit(job_info, status)
+
+    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]:
+        """
+        Creates and polls a remote inference job asynchronously.
+        Reuses existing synchronous methods but runs them in an async context.
+
+        :param iterations: Number of times to run each interview
+        :param remote_inference_description: Optional description for the remote job
+        :param remote_inference_results_visibility: Visibility setting for results
+        :return: Results object if successful, None if job fails or is cancelled
+        """
+        import asyncio
+        from functools import partial
+
+        # Create job using existing method
+        loop = asyncio.get_event_loop()
+        job_info = await loop.run_in_executor(
+            None,
+            partial(
+                self.create_remote_inference_job,
+                iterations=iterations,
+                remote_inference_description=remote_inference_description,
+                remote_inference_results_visibility=remote_inference_results_visibility,
+            ),
+        )
+        if job_info is None:
+            raise ValueError("Remote job creation failed.")
+
+        return await loop.run_in_executor(
+            None,
+            partial(self.poll_remote_inference_job, job_info),
+        )
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)