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

edsl/jobs/JobsRemoteInferenceLogger.py

@@ -0,0 +1,239 @@
+import re
+import sys
+import uuid
+from abc import ABC, abstractmethod
+from typing import Optional, Union, Literal, TYPE_CHECKING, List, Dict
+from datetime import datetime
+from dataclasses import dataclass
+from edsl.exceptions.coop import CoopServerResponseError
+
+from edsl.jobs.jobs_status_enums import JobsStatus
+
+if TYPE_CHECKING:
+    from edsl.results.Results import Results
+
+
+@dataclass
+class LogMessage:
+    text: str
+    status: str
+    timestamp: datetime
+    status: JobsStatus
+
+
+@dataclass
+class JobsInfo:
+    job_uuid: str = None
+    progress_bar_url: str = None
+    error_report_url: str = None
+    results_uuid: str = None
+    results_url: str = None
+
+    pretty_names = {
+        "job_uuid": "Job UUID",
+        "progress_bar_url": "Progress Bar URL",
+        "error_report_url": "Error Report URL",
+        "results_uuid": "Results UUID",
+        "results_url": "Results URL",
+    }
+
+
+class JobLogger(ABC):
+    def __init__(self, verbose: bool = False):
+        self.verbose = verbose
+        self.jobs_info = JobsInfo()
+
+    def add_info(
+        self,
+        information_type: Literal[
+            "job_uuid",
+            "progress_bar_url",
+            "error_report_url",
+            "results_uuid",
+            "results_url",
+        ],
+        value: str,
+    ):
+        """Add information to the logger
+
+        >>> j = StdOutJobLogger()
+        >>> j.add_info("job_uuid", "1234")
+        >>> j.jobs_info.job_uuid
+        '1234'
+        """
+        if information_type not in self.jobs_info.__annotations__:
+            raise ValueError(f"Information type {information_type} not supported")
+        setattr(self.jobs_info, information_type, value)
+
+    @abstractmethod
+    def update(self, message: str, status: str = "running"):
+        pass
+
+
+class HTMLTableJobLogger(JobLogger):
+    def __init__(self, verbose=True, **kwargs):
+        from IPython.display import display, HTML
+
+        super().__init__(verbose=verbose)
+        self.display_handle = display(HTML(""), display_id=True)
+        self.current_message = None
+        self.log_id = str(uuid.uuid4())
+        self.is_expanded = True
+        self.spinner_chars = ["◐", "◓", "◑", "◒"]  # Rotating spinner characters
+        self.spinner_idx = 0
+
+    def _get_table_row(self, key: str, value: str) -> str:
+        """Generate a table row with key-value pair"""
+        return f"""
+            <tr>
+                <td style="padding: 8px; border: 1px solid #ddd; font-weight: bold;">{key}</td>
+                <td style="padding: 8px; border: 1px solid #ddd;">{value if value else 'None'}</td>
+            </tr>
+        """
+
+    def _linkify(self, text: str) -> str:
+        """Convert URLs in text to clickable links"""
+        url_pattern = r'(https?://[^\s<>"]+|www\.[^\s<>"]+)'
+        return re.sub(
+            url_pattern,
+            r'<a href="\1" target="_blank" style="color: #3b82f6; text-decoration: underline;">\1</a>',
+            text,
+        )
+
+    def _get_spinner(self, status: JobsStatus) -> str:
+        """Get the current spinner frame if status is running"""
+        if status == JobsStatus.RUNNING:
+            spinner = self.spinner_chars[self.spinner_idx]
+            self.spinner_idx = (self.spinner_idx + 1) % len(self.spinner_chars)
+            return f'<span style="margin-right: 8px;">{spinner}</span>'
+        elif status == JobsStatus.COMPLETED:
+            return '<span style="margin-right: 8px; color: #22c55e;">✓</span>'
+        elif status == JobsStatus.FAILED:
+            return '<span style="margin-right: 8px; color: #ef4444;">✗</span>'
+        return ""
+
+    def _get_html(self, status: JobsStatus = JobsStatus.RUNNING) -> str:
+        """Generate the complete HTML display"""
+        # Generate table rows for each JobsInfo field
+        info_rows = ""
+        for field, _ in self.jobs_info.__annotations__.items():
+            if field != "pretty_names":  # Skip the pretty_names dictionary
+                value = getattr(self.jobs_info, field)
+                value = self._linkify(str(value)) if value else None
+                pretty_name = self.jobs_info.pretty_names.get(
+                    field, field.replace("_", " ").title()
+                )
+                info_rows += self._get_table_row(pretty_name, value)
+
+        # Add current message section with spinner
+        message_html = ""
+        if self.current_message:
+            spinner = self._get_spinner(status)
+            message_html = f"""
+                <div style="margin-top: 10px; padding: 8px; background-color: #f8f9fa; border: 1px solid #ddd; border-radius: 4px;">
+                    {spinner}<strong>Current Status:</strong> {self._linkify(self.current_message)}
+                </div>
+            """
+
+        display_style = "block" if self.is_expanded else "none"
+        arrow = "▼" if self.is_expanded else "▶"
+
+        return f"""
+            <div style="font-family: system-ui; max-width: 800px; margin: 10px 0;">
+                <div onclick="document.getElementById('content-{self.log_id}').style.display = document.getElementById('content-{self.log_id}').style.display === 'none' ? 'block' : 'none';
+                             document.getElementById('arrow-{self.log_id}').innerHTML = document.getElementById('content-{self.log_id}').style.display === 'none' ? '▶' : '▼';"
+                     style="padding: 10px; background: #f5f5f5; border: 1px solid #ddd; border-radius: 4px; cursor: pointer;">
+                    <span id="arrow-{self.log_id}">{arrow}</span> Job Status ({datetime.now().strftime('%Y-%m-%d %H:%M:%S')})
+                </div>
+                <div id="content-{self.log_id}" style="display: {display_style};">
+                    <table style="width: 100%; border-collapse: collapse; background: white; border: 1px solid #ddd;">
+                        {info_rows}
+                    </table>
+                    {message_html}
+                </div>
+            </div>
+        """
+
+    def update(self, message: str, status: JobsStatus = JobsStatus.RUNNING):
+        """Update the display with new message and current JobsInfo state"""
+        from IPython.display import HTML
+
+        self.current_message = message
+        if self.verbose:
+            self.display_handle.update(HTML(self._get_html(status)))
+        else:
+            return None
+
+
+class StdOutJobLogger(JobLogger):
+    def __init__(self, verbose=True, **kwargs):
+        super().__init__(verbose=verbose)  # Properly call parent's __init__
+        self.messages: List[LogMessage] = []
+
+    def update(self, message: str, status: JobsStatus = JobsStatus.RUNNING):
+        log_msg = LogMessage(text=message, status=status, timestamp=datetime.now())
+        self.messages.append(log_msg)
+        if self.verbose:
+            sys.stdout.write(f"│ {message}\n")
+            sys.stdout.flush()
+        else:
+            return None
+
+
+class JupyterJobLogger(JobLogger):
+    def __init__(self, verbose=True, **kwargs):
+        from IPython.display import display, HTML
+
+        super().__init__(verbose=verbose)
+        self.messages = []
+        self.log_id = str(uuid.uuid4())
+        self.is_expanded = True
+        self.display_handle = display(HTML(""), display_id=True)
+
+    def _linkify(self, text):
+        url_pattern = r'(https?://[^\s<>"]+|www\.[^\s<>"]+)'
+        return re.sub(
+            url_pattern,
+            r'<a href="\1" target="_blank" style="color: #3b82f6; text-decoration: underline;">\1</a>',
+            text,
+        )
+
+    def _get_html(self):
+        messages_html = "\n".join(
+            [
+                f'<div style="border-left: 3px solid {msg["color"]}; padding: 5px 10px; margin: 5px 0;">{self._linkify(msg["text"])}</div>'
+                for msg in self.messages
+            ]
+        )
+
+        display_style = "block" if self.is_expanded else "none"
+        arrow = "▼" if self.is_expanded else "▶"
+
+        return f"""
+            <div style="border: 1px solid #ccc; margin: 10px 0; max-width: 800px;">
+                <div onclick="document.getElementById('content-{self.log_id}').style.display = document.getElementById('content-{self.log_id}').style.display === 'none' ? 'block' : 'none';
+                             document.getElementById('arrow-{self.log_id}').innerHTML = document.getElementById('content-{self.log_id}').style.display === 'none' ? '▶' : '▼';"
+                     style="padding: 10px; background: #f5f5f5; cursor: pointer;">
+                    <span id="arrow-{self.log_id}">{arrow}</span> Remote Job Log ({datetime.now().strftime('%Y-%m-%d %H:%M:%S')})
+                </div>
+                <div id="content-{self.log_id}" style="padding: 10px; display: {display_style};">
+                    {messages_html}
+                </div>
+            </div>
+        """
+
+    def update(self, message, status: JobsStatus = JobsStatus.RUNNING):
+        from IPython.display import HTML
+
+        colors = {"running": "#3b82f6", "completed": "#22c55e", "failed": "#ef4444"}
+        self.messages.append({"text": message, "color": colors.get(status, "#666")})
+        if self.verbose:
+            self.display_handle.update(HTML(self._get_html()))
+        else:
+            return None
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/jobs/RequestTokenEstimator.py

@@ -0,0 +1,30 @@
+from edsl.jobs.FetchInvigilator import FetchInvigilator
+
+
+class RequestTokenEstimator:
+    """Estimate the number of tokens that will be required to run the focal task."""
+
+    def __init__(self, interview):
+        self.interview = interview
+
+    def __call__(self, question) -> float:
+        """Estimate the number of tokens that will be required to run the focal task."""
+        from edsl.scenarios.FileStore import FileStore
+
+        invigilator = FetchInvigilator(self.interview)(question=question)
+
+        # TODO: There should be a way to get a more accurate estimate.
+        combined_text = ""
+        file_tokens = 0
+        for prompt in invigilator.get_prompts().values():
+            if hasattr(prompt, "text"):
+                combined_text += prompt.text
+            elif isinstance(prompt, str):
+                combined_text += prompt
+            elif isinstance(prompt, list):
+                for file in prompt:
+                    if isinstance(file, FileStore):
+                        file_tokens += file.size * 0.25
+            else:
+                raise ValueError(f"Prompt is of type {type(prompt)}")
+        return len(combined_text) / 4.0 + file_tokens

edsl/jobs/__init__.py

@@ -1 +1 @@
-from edsl.jobs.Jobs import Jobs
+from edsl.jobs.Jobs import Jobs

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/buckets/BucketCollection.py

@@ -1,63 +1,104 @@
-from collections import UserDict
-from edsl.jobs.buckets.TokenBucket import TokenBucket
-from edsl.jobs.buckets.ModelBuckets import ModelBuckets
-
-
-class BucketCollection(UserDict):
-    """A Jobs object will have a whole collection of model buckets, as multiple models could be used.
-
-    The keys here are the models, and the values are the ModelBuckets objects.
-    Models themselves are hashable, so this works.
-    """
-
-    def __init__(self, infinity_buckets=False):
-        super().__init__()
-        self.infinity_buckets = infinity_buckets
-        self.models_to_services = {}
-        self.services_to_buckets = {}
-
-    def __repr__(self):
-        return f"BucketCollection({self.data})"
-
-    def add_model(self, model: "LanguageModel") -> None:
-        """Adds a model to the bucket collection.
-
-        This will create the token and request buckets for the model."""
-
-        # compute the TPS and RPS from the model
-        if not self.infinity_buckets:
-            TPS = model.TPM / 60.0
-            RPS = model.RPM / 60.0
-        else:
-            TPS = float("inf")
-            RPS = float("inf")
-
-        if model.model not in self.models_to_services:
-            service = model._inference_service_
-            if service not in self.services_to_buckets:
-                requests_bucket = TokenBucket(
-                    bucket_name=service,
-                    bucket_type="requests",
-                    capacity=RPS,
-                    refill_rate=RPS,
-                )
-                tokens_bucket = TokenBucket(
-                    bucket_name=service,
-                    bucket_type="tokens",
-                    capacity=TPS,
-                    refill_rate=TPS,
-                )
-                self.services_to_buckets[service] = ModelBuckets(
-                    requests_bucket, tokens_bucket
-                )
-            self.models_to_services[model.model] = service
-            self[model] = self.services_to_buckets[service]
-        else:
-            self[model] = self.services_to_buckets[self.models_to_services[model.model]]
-
-    def visualize(self) -> dict:
-        """Visualize the token and request buckets for each model."""
-        plots = {}
-        for model in self:
-            plots[model] = self[model].visualize()
-        return plots
+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.
+
+    The keys here are the models, and the values are the ModelBuckets objects.
+    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.
+        """
+        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})"
+
+    def add_model(self, model: "LanguageModel") -> None:
+        """Adds a model to the bucket collection.
+
+        This will create the token and request buckets for the model."""
+
+        # compute the TPS and RPS from the model
+        if not self.infinity_buckets:
+            TPS = model.tpm / 60.0
+            RPS = model.rpm / 60.0
+        else:
+            TPS = float("inf")
+            RPS = float("inf")
+
+        if model.model not in self.models_to_services:
+            service = model._inference_service_
+            if service not in self.services_to_buckets:
+                requests_bucket = TokenBucket(
+                    bucket_name=service,
+                    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
+                )
+            self.models_to_services[model.model] = service
+            self[model] = self.services_to_buckets[service]
+        else:
+            self[model] = self.services_to_buckets[self.models_to_services[model.model]]
+
+    def visualize(self) -> dict:
+        """Visualize the token and request buckets for each model."""
+        plots = {}
+        for model in self:
+            plots[model] = self[model].visualize()
+        return plots