edsl 0.1.15__py3-none-any.whl → 0.1.40__py3-none-any.whl

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

@@ -0,0 +1,104 @@
+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

edsl/jobs/buckets/ModelBuckets.py

@@ -0,0 +1,65 @@
+# from edsl.jobs.buckets.TokenBucket import TokenBucket
+
+
+class ModelBuckets:
+    """A class to represent the token and request buckets for a model.
+
+    Most LLM model services have limits both on requests-per-minute (RPM) and tokens-per-minute (TPM).
+    A request is one call to the service. The number of tokens required for a request depends on parameters.
+    """
+
+    def __init__(self, requests_bucket: "TokenBucket", tokens_bucket: "TokenBucket"):
+        """Initialize the model buckets.
+
+        The requests bucket captures requests per unit of time.
+        The tokens bucket captures the number of language model tokens.
+
+        """
+        self.requests_bucket = requests_bucket
+        self.tokens_bucket = tokens_bucket
+
+    def __add__(self, other: "ModelBuckets"):
+        """Combine two model buckets."""
+        return ModelBuckets(
+            requests_bucket=self.requests_bucket + other.requests_bucket,
+            tokens_bucket=self.tokens_bucket + other.tokens_bucket,
+        )
+
+    def turbo_mode_on(self):
+        """Set the refill rate to infinity for both buckets."""
+        self.requests_bucket.turbo_mode_on()
+        self.tokens_bucket.turbo_mode_on()
+
+    def turbo_mode_off(self):
+        """Restore the refill rate to its original value for both buckets."""
+        self.requests_bucket.turbo_mode_off()
+        self.tokens_bucket.turbo_mode_off()
+
+    @classmethod
+    def infinity_bucket(cls, model_name: str = "not_specified") -> "ModelBuckets":
+        """Create a bucket with infinite capacity and refill rate."""
+        from edsl.jobs.buckets.TokenBucket import TokenBucket
+
+        return cls(
+            requests_bucket=TokenBucket(
+                bucket_name=model_name,
+                bucket_type="requests",
+                capacity=float("inf"),
+                refill_rate=float("inf"),
+            ),
+            tokens_bucket=TokenBucket(
+                bucket_name=model_name,
+                bucket_type="tokens",
+                capacity=float("inf"),
+                refill_rate=float("inf"),
+            ),
+        )
+
+    def visualize(self):
+        """Visualize the token and request buckets."""
+        plot1 = self.requests_bucket.visualize()
+        plot2 = self.tokens_bucket.visualize()
+        return plot1, plot2
+
+    def __repr__(self):
+        return f"ModelBuckets(requests_bucket={self.requests_bucket}, tokens_bucket={self.tokens_bucket})"

edsl/jobs/buckets/TokenBucket.py

@@ -0,0 +1,283 @@
+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.
+    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,
+        *,
+        bucket_name: str,
+        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
+        self.added_tokens = 0
+        self._lock = RLock()
+
+        self.target_rate = (
+            capacity * 60
+        )  # set this here because it can change with turbo mode
+
+        self._old_capacity = capacity
+        self.tokens = capacity  # Current number of available tokens
+        self.refill_rate = refill_rate  # Rate at which tokens are refilled
+        self._old_refill_rate = refill_rate
+        self.last_refill = time.monotonic()  # Last refill time
+        self.log: List[Any] = []
+        self.turbo_mode = False
+
+        self.creation_time = time.monotonic()
+
+        self.num_requests = 0
+        self.num_released = 0
+        self.tokens_returned = 0
+
+    def turbo_mode_on(self):
+        """Set the refill rate to infinity."""
+        if self.turbo_mode:
+            pass
+        else:
+            # pass
+            self.turbo_mode = True
+            self.capacity = float("inf")
+            self.refill_rate = float("inf")
+
+    def turbo_mode_off(self):
+        """Restore the refill rate to its original value."""
+        self.turbo_mode = False
+        self.capacity = self._old_capacity
+        self.refill_rate = self._old_refill_rate
+
+    def __add__(self, other) -> "TokenBucket":
+        """Combine two token buckets.
+
+        The resulting bucket has the minimum capacity and refill rate of the two buckets.
+        This is useful, for example, if we have two calls to the same model on the same service but have different temperatures.
+        """
+        return TokenBucket(
+            bucket_name=self.bucket_name,
+            bucket_type=self.bucket_type,
+            capacity=min(self.capacity, other.capacity),
+            refill_rate=min(self.refill_rate, other.refill_rate),
+        )
+
+    def __repr__(self):
+        return f"TokenBucket(bucket_name={self.bucket_name}, bucket_type='{self.bucket_type}', capacity={self.capacity}, refill_rate={self.refill_rate})"
+
+    def add_tokens(self, tokens: Union[int, float]) -> None:
+        """Add tokens to the bucket, up to the maximum capacity.
+
+        :param tokens: The number of tokens to add to the bucket.
+
+        >>> bucket = TokenBucket(bucket_name="test", bucket_type="test", capacity=10, refill_rate=1)
+        >>> bucket.tokens
+        10
+        >>> bucket.add_tokens(5)
+        >>> bucket.tokens
+        10
+        """
+        self.tokens_returned += tokens
+        self.tokens = min(self.capacity, self.tokens + tokens)
+        self.log.append((time.monotonic(), self.tokens))
+
+    def refill(self) -> None:
+        """Refill the bucket with new tokens based on elapsed time.
+
+
+
+        >>> bucket = TokenBucket(bucket_name="test", bucket_type="test", capacity=10, refill_rate=1)
+        >>> bucket.tokens = 0
+        >>> bucket.refill()
+        >>> bucket.tokens > 0
+        True
+        """
+        """Refill the bucket with new tokens based on elapsed time."""
+        now = time.monotonic()
+        # print(f"Time is now: {now}; Last refill time: {self.last_refill}")
+        elapsed = now - self.last_refill
+        # print("Elapsed time: ", elapsed)
+        refill_amount = elapsed * self.refill_rate
+        self.tokens = min(self.capacity, self.tokens + refill_amount)
+        self.last_refill = now
+
+        if self.tokens < self.capacity:
+            pass
+            # print(f"Refilled. Current tokens: {self.tokens:.4f}")
+            # print(f"Elapsed time: {elapsed:.4f} seconds")
+            # print(f"Refill amount: {refill_amount:.4f}")
+
+        self.log.append((now, self.tokens))
+
+    def wait_time(self, requested_tokens: Union[float, int]) -> float:
+        """Calculate the time to wait for the requested number of tokens."""
+        # self.refill()  # Update the current token count
+        if self.tokens >= requested_tokens:
+            return 0
+        return (requested_tokens - self.tokens) / self.refill_rate
+
+    async def get_tokens(
+        self, amount: Union[int, float] = 1, cheat_bucket_capacity=True
+    ) -> None:
+        """Wait for the specified number of tokens to become available.
+
+
+        :param amount: The number of tokens
+        :param warn: If True, warn if the requested amount exceeds the bucket capacity.
+
+        >>> bucket = TokenBucket(bucket_name="test", bucket_type="test", capacity=10, refill_rate=1)
+        >>> asyncio.run(bucket.get_tokens(5))
+        >>> bucket.tokens
+        5
+        >>> asyncio.run(bucket.get_tokens(9))
+        >>> bucket.tokens < 1
+        True
+
+        >>> bucket = TokenBucket(bucket_name="test", bucket_type="test", capacity=10, refill_rate=1)
+        >>> asyncio.run(bucket.get_tokens(11, cheat_bucket_capacity=False))
+        Traceback (most recent call last):
+        ...
+        ValueError: Requested amount exceeds bucket capacity. Bucket capacity: 10, requested amount: 11. As the bucket never overflows, the requested amount will never be available.
+        >>> asyncio.run(bucket.get_tokens(11, cheat_bucket_capacity=True))
+        >>> bucket.capacity
+        12.100000000000001
+        """
+        self.num_requests += amount
+        if amount >= self.capacity:
+            if not cheat_bucket_capacity:
+                msg = f"Requested amount exceeds bucket capacity. Bucket capacity: {self.capacity}, requested amount: {amount}. As the bucket never overflows, the requested amount will never be available."
+                raise ValueError(msg)
+            else:
+                self.capacity = amount * 1.10
+                self._old_capacity = self.capacity
+
+        start_time = time.monotonic()
+        while True:
+            self.refill()  # Refill based on elapsed time
+            if self.tokens >= amount:
+                self.tokens -= amount
+                break
+
+            wait_time = self.wait_time(amount)
+            if wait_time > 0:
+                await asyncio.sleep(wait_time)
+
+        self.num_released += amount
+        now = time.monotonic()
+        self.log.append((now, self.tokens))
+        return None
+
+    def get_log(self) -> list[tuple]:
+        return self.log
+
+    def visualize(self):
+        """Visualize the token bucket over time."""
+        times, tokens = zip(*self.get_log())
+        start_time = times[0]
+        times = [t - start_time for t in times]  # Normalize time to start from 0
+        from matplotlib import pyplot as plt
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(times, tokens, label="Tokens Available")
+        plt.xlabel("Time (seconds)", fontsize=12)
+        plt.ylabel("Number of Tokens", fontsize=12)
+        details = f"{self.bucket_name} ({self.bucket_type}) Bucket Usage Over Time\nCapacity: {self.capacity:.1f}, Refill Rate: {self.refill_rate:.1f}/second"
+        plt.title(details, fontsize=14)
+
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+        plt.show()
+
+    def get_throughput(self, time_window: Optional[float] = None) -> float:
+        """
+        Calculate the empirical bucket throughput in tokens per minute for the specified time window.
+
+        :param time_window: The time window in seconds to calculate the throughput for.
+        :return: The throughput in tokens per minute.
+
+        >>> bucket = TokenBucket(bucket_name="test", bucket_type="test", capacity=100, refill_rate=10)
+        >>> asyncio.run(bucket.get_tokens(50))
+        >>> time.sleep(1)  # Wait for 1 second
+        >>> asyncio.run(bucket.get_tokens(30))
+        >>> throughput = bucket.get_throughput(1)
+        >>> 4750 < throughput < 4850
+        True
+        """
+        now = time.monotonic()
+
+        if time_window is None:
+            start_time = self.creation_time
+        else:
+            start_time = now - time_window
+
+        if start_time < self.creation_time:
+            start_time = self.creation_time
+
+        elapsed_time = now - start_time
+
+        if elapsed_time == 0:
+            return self.num_released / 0.001
+
+        return (self.num_released / elapsed_time) * 60
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)