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

edsl/jobs/buckets/BucketCollection.py

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

edsl/jobs/buckets/ModelBuckets.py

@@ -1,65 +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})"
+# 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})"