edsl 0.1.39.dev1__py3-none-any.whl → 0.1.39.dev2__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/buckets/BucketCollection.py

@@ -1,8 +1,15 @@
+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.
 
@@ -10,11 +17,43 @@ class BucketCollection(UserDict):
     Models themselves are hashable, so this works.
     """
 
-    def __init__(self, infinity_buckets=False):
+    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})"
@@ -26,8 +65,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")
@@ -40,12 +79,14 @@ 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

edsl/jobs/buckets/TokenBucket.py

@@ -1,10 +1,55 @@
 from typing import Union, List, Any, Optional
 import asyncio
 import time
+from threading import RLock
+from edsl.jobs.decorators import synchronized_class
 
+from typing import Union, List, Any, Optional
+import asyncio
+import time
+from threading import RLock
+from edsl.jobs.decorators import synchronized_class
 
+
+@synchronized_class
 class TokenBucket:
-    """This is a token bucket used to respect rate limits to services."""
+    """This is a token bucket used to respect rate limits to services.
+    It can operate either locally or remotely via a REST API based on initialization parameters.
+    """
+
+    def __new__(
+        cls,
+        *,
+        bucket_name: str,
+        bucket_type: str,
+        capacity: Union[int, float],
+        refill_rate: Union[int, float],
+        remote_url: Optional[str] = None,
+    ):
+        """Factory method to create either a local or remote token bucket.
+
+        Args:
+            bucket_name: Name of the bucket
+            bucket_type: Type of the bucket
+            capacity: Maximum number of tokens
+            refill_rate: Rate at which tokens are refilled
+            remote_url: If provided, creates a remote token bucket client
+        """
+        if remote_url is not None:
+            # Import here to avoid circular imports
+            from edsl.jobs.buckets.TokenBucketClient import TokenBucketClient
+
+            return TokenBucketClient(
+                bucket_name=bucket_name,
+                bucket_type=bucket_type,
+                capacity=capacity,
+                refill_rate=refill_rate,
+                api_base_url=remote_url,
+            )
+
+        # Create a local token bucket
+        instance = super(TokenBucket, cls).__new__(cls)
+        return instance
 
     def __init__(
         self,
@@ -13,11 +58,17 @@ class TokenBucket:
         bucket_type: str,
         capacity: Union[int, float],
         refill_rate: Union[int, float],
+        remote_url: Optional[str] = None,
     ):
+        # Skip initialization if this is a remote bucket
+        if remote_url is not None:
+            return
+
         self.bucket_name = bucket_name
         self.bucket_type = bucket_type
-        self.capacity = capacity  # Maximum number of tokens
+        self.capacity = capacity
         self.added_tokens = 0
+        self._lock = RLock()
 
         self.target_rate = (
             capacity * 60
@@ -225,25 +276,6 @@ class TokenBucket:
 
         return (self.num_released / elapsed_time) * 60
 
-        # # Filter log entries within the time window
-        # relevant_log = [(t, tokens) for t, tokens in self.log if t >= start_time]
-
-        # if len(relevant_log) < 2:
-        #     return 0  # Not enough data points to calculate throughput
-
-        # # Calculate total tokens used
-        # initial_tokens = relevant_log[0][1]
-        # final_tokens = relevant_log[-1][1]
-        # tokens_used = self.num_released - (final_tokens - initial_tokens)
-
-        # # Calculate actual time elapsed
-        # actual_time_elapsed = relevant_log[-1][0] - relevant_log[0][0]
-
-        # # Calculate throughput in tokens per minute
-        # throughput = (tokens_used / actual_time_elapsed) * 60
-
-        # return throughput
-
 
 if __name__ == "__main__":
     import doctest

edsl/jobs/buckets/TokenBucketAPI.py

@@ -0,0 +1,211 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Union, Dict
+from typing import Union, List, Any, Optional
+from threading import RLock
+from edsl.jobs.buckets.TokenBucket import TokenBucket  # Original implementation
+
+
+def safe_float_for_json(value: float) -> Union[float, str]:
+    """Convert float('inf') to 'infinity' for JSON serialization.
+
+    Args:
+        value: The float value to convert
+
+    Returns:
+        Either the original float or the string 'infinity' if the value is infinite
+    """
+    if value == float("inf"):
+        return "infinity"
+    return value
+
+
+app = FastAPI()
+
+# In-memory storage for TokenBucket instances
+buckets: Dict[str, TokenBucket] = {}
+
+
+class TokenBucketCreate(BaseModel):
+    bucket_name: str
+    bucket_type: str
+    capacity: Union[int, float]
+    refill_rate: Union[int, float]
+
+
+@app.get("/buckets")
+async def list_buckets(
+    bucket_type: Optional[str] = None,
+    bucket_name: Optional[str] = None,
+    include_logs: bool = False,
+):
+    """List all buckets and their current status.
+
+    Args:
+        bucket_type: Optional filter by bucket type
+        bucket_name: Optional filter by bucket name
+        include_logs: Whether to include the full logs in the response
+    """
+    result = {}
+
+    for bucket_id, bucket in buckets.items():
+        # Apply filters if specified
+        if bucket_type and bucket.bucket_type != bucket_type:
+            continue
+        if bucket_name and bucket.bucket_name != bucket_name:
+            continue
+
+        # Get basic bucket info
+        bucket_info = {
+            "bucket_name": bucket.bucket_name,
+            "bucket_type": bucket.bucket_type,
+            "tokens": bucket.tokens,
+            "capacity": bucket.capacity,
+            "refill_rate": bucket.refill_rate,
+            "turbo_mode": bucket.turbo_mode,
+            "num_requests": bucket.num_requests,
+            "num_released": bucket.num_released,
+            "tokens_returned": bucket.tokens_returned,
+        }
+        for k, v in bucket_info.items():
+            if isinstance(v, float):
+                bucket_info[k] = safe_float_for_json(v)
+
+        # Only include logs if requested
+        if include_logs:
+            bucket_info["log"] = bucket.log
+
+        result[bucket_id] = bucket_info
+
+    return result
+
+
+@app.post("/bucket/{bucket_id}/add_tokens")
+async def add_tokens(bucket_id: str, amount: float):
+    """Add tokens to an existing bucket."""
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    if not isinstance(amount, (int, float)) or amount != amount:  # Check for NaN
+        raise HTTPException(status_code=400, detail="Invalid amount specified")
+
+    if amount == float("inf") or amount == float("-inf"):
+        raise HTTPException(status_code=400, detail="Amount cannot be infinite")
+
+    bucket = buckets[bucket_id]
+    bucket.add_tokens(amount)
+
+    # Ensure we return a JSON-serializable float
+    current_tokens = float(bucket.tokens)
+    if not -1e308 <= current_tokens <= 1e308:  # Check if within JSON float bounds
+        current_tokens = 0.0  # or some other reasonable default
+
+    return {"status": "success", "current_tokens": safe_float_for_json(current_tokens)}
+
+
+# @app.post("/bucket")
+# async def create_bucket(bucket: TokenBucketCreate):
+#     bucket_id = f"{bucket.bucket_name}_{bucket.bucket_type}"
+#     if bucket_id in buckets:
+#         raise HTTPException(status_code=400, detail="Bucket already exists")
+
+#     # Create an actual TokenBucket instance
+#     buckets[bucket_id] = TokenBucket(
+#         bucket_name=bucket.bucket_name,
+#         bucket_type=bucket.bucket_type,
+#         capacity=bucket.capacity,
+#         refill_rate=bucket.refill_rate,
+#     )
+#     return {"status": "created"}
+
+
+@app.post("/bucket")
+async def create_bucket(bucket: TokenBucketCreate):
+    if (
+        not isinstance(bucket.capacity, (int, float))
+        or bucket.capacity != bucket.capacity
+    ):  # Check for NaN
+        raise HTTPException(status_code=400, detail="Invalid capacity value")
+    if (
+        not isinstance(bucket.refill_rate, (int, float))
+        or bucket.refill_rate != bucket.refill_rate
+    ):  # Check for NaN
+        raise HTTPException(status_code=400, detail="Invalid refill rate value")
+    if bucket.capacity == float("inf") or bucket.refill_rate == float("inf"):
+        raise HTTPException(status_code=400, detail="Values cannot be infinite")
+    bucket_id = f"{bucket.bucket_name}_{bucket.bucket_type}"
+    if bucket_id in buckets:
+        # Instead of error, return success with "existing" status
+        return {
+            "status": "existing",
+            "bucket": {
+                "capacity": safe_float_for_json(buckets[bucket_id].capacity),
+                "refill_rate": safe_float_for_json(buckets[bucket_id].refill_rate),
+            },
+        }
+
+    # Create a new bucket
+    buckets[bucket_id] = TokenBucket(
+        bucket_name=bucket.bucket_name,
+        bucket_type=bucket.bucket_type,
+        capacity=bucket.capacity,
+        refill_rate=bucket.refill_rate,
+    )
+    return {"status": "created"}
+
+
+@app.post("/bucket/{bucket_id}/get_tokens")
+async def get_tokens(bucket_id: str, amount: float, cheat_bucket_capacity: bool = True):
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    bucket = buckets[bucket_id]
+    await bucket.get_tokens(amount, cheat_bucket_capacity)
+    return {"status": "success"}
+
+
+@app.post("/bucket/{bucket_id}/turbo_mode/{state}")
+async def set_turbo_mode(bucket_id: str, state: bool):
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    bucket = buckets[bucket_id]
+    if state:
+        bucket.turbo_mode_on()
+    else:
+        bucket.turbo_mode_off()
+    return {"status": "success"}
+
+
+@app.get("/bucket/{bucket_id}/status")
+async def get_bucket_status(bucket_id: str):
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    bucket = buckets[bucket_id]
+    status = {
+        "tokens": bucket.tokens,
+        "capacity": bucket.capacity,
+        "refill_rate": bucket.refill_rate,
+        "turbo_mode": bucket.turbo_mode,
+        "num_requests": bucket.num_requests,
+        "num_released": bucket.num_released,
+        "tokens_returned": bucket.tokens_returned,
+        "log": bucket.log,
+    }
+    for k, v in status.items():
+        if isinstance(v, float):
+            status[k] = safe_float_for_json(v)
+
+    for index, entry in enumerate(status["log"]):
+        ts, value = entry
+        status["log"][index] = (ts, safe_float_for_json(value))
+
+    # print(status)
+    return status
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8001)