edsl 0.1.38.dev4__py3-none-any.whl → 0.1.39__py3-none-any.whl

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/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,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)