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

edsl/jobs/buckets/TokenBucketAPI.py

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

edsl/jobs/buckets/TokenBucketClient.py

@@ -1,191 +0,0 @@
-from typing import Union, Optional
-import asyncio
-import time
-import aiohttp
-
-
-class TokenBucketClient:
-    """REST API client version of TokenBucket that maintains the same interface
-    by delegating to a server running the original TokenBucket implementation."""
-
-    def __init__(
-        self,
-        *,
-        bucket_name: str,
-        bucket_type: str,
-        capacity: Union[int, float],
-        refill_rate: Union[int, float],
-        api_base_url: str = "http://localhost:8000",
-    ):
-        self.bucket_name = bucket_name
-        self.bucket_type = bucket_type
-        self.capacity = capacity
-        self.refill_rate = refill_rate
-        self.api_base_url = api_base_url
-        self.bucket_id = f"{bucket_name}_{bucket_type}"
-
-        # Initialize the bucket on the server
-        asyncio.run(self._create_bucket())
-
-        # Cache some values locally
-        self.creation_time = time.monotonic()
-        self.turbo_mode = False
-
-    async def _create_bucket(self):
-        async with aiohttp.ClientSession() as session:
-            payload = {
-                "bucket_name": self.bucket_name,
-                "bucket_type": self.bucket_type,
-                "capacity": self.capacity,
-                "refill_rate": self.refill_rate,
-            }
-            async with session.post(
-                f"{self.api_base_url}/bucket",
-                json=payload,
-            ) as response:
-                if response.status != 200:
-                    raise ValueError(f"Unexpected error: {await response.text()}")
-
-                result = await response.json()
-                if result["status"] == "existing":
-                    # Update our local values to match the existing bucket
-                    self.capacity = float(result["bucket"]["capacity"])
-                    self.refill_rate = float(result["bucket"]["refill_rate"])
-
-    def turbo_mode_on(self):
-        """Set the refill rate to infinity."""
-        asyncio.run(self._set_turbo_mode(True))
-        self.turbo_mode = True
-
-    def turbo_mode_off(self):
-        """Restore the refill rate to its original value."""
-        asyncio.run(self._set_turbo_mode(False))
-        self.turbo_mode = False
-
-    async def add_tokens(self, amount: Union[int, float]):
-        """Add tokens to the bucket."""
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                f"{self.api_base_url}/bucket/{self.bucket_id}/add_tokens",
-                params={"amount": amount},
-            ) as response:
-                if response.status != 200:
-                    raise ValueError(f"Failed to add tokens: {await response.text()}")
-
-    async def _set_turbo_mode(self, state: bool):
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                f"{self.api_base_url}/bucket/{self.bucket_id}/turbo_mode/{str(state).lower()}"
-            ) as response:
-                if response.status != 200:
-                    raise ValueError(
-                        f"Failed to set turbo mode: {await response.text()}"
-                    )
-
-    async def get_tokens(
-        self, amount: Union[int, float] = 1, cheat_bucket_capacity=True
-    ) -> None:
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                f"{self.api_base_url}/bucket/{self.bucket_id}/get_tokens",
-                params={
-                    "amount": amount,
-                    "cheat_bucket_capacity": int(cheat_bucket_capacity),
-                },
-            ) as response:
-                if response.status != 200:
-                    raise ValueError(f"Failed to get tokens: {await response.text()}")
-
-    def get_throughput(self, time_window: Optional[float] = None) -> float:
-        status = asyncio.run(self._get_status())
-        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 status["num_released"] / 0.001
-
-        return (status["num_released"] / elapsed_time) * 60
-
-    async def _get_status(self) -> dict:
-        async with aiohttp.ClientSession() as session:
-            async with session.get(
-                f"{self.api_base_url}/bucket/{self.bucket_id}/status"
-            ) as response:
-                if response.status != 200:
-                    raise ValueError(
-                        f"Failed to get bucket status: {await response.text()}"
-                    )
-                return await response.json()
-
-    def __add__(self, other) -> "TokenBucketClient":
-        """Combine two token buckets."""
-        return TokenBucketClient(
-            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),
-            api_base_url=self.api_base_url,
-        )
-
-    @property
-    def tokens(self) -> float:
-        """Get the number of tokens remaining in the bucket."""
-        status = asyncio.run(self._get_status())
-        return float(status["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 >= float(requested_tokens):
-            return 0.0
-        try:
-            return (requested_tokens - self.tokens) / self.refill_rate
-        except Exception as e:
-            raise ValueError(f"Error calculating wait time: {e}")
-
-    # def wait_time(self, num_tokens: Union[int, float]) -> float:
-    #     return 0  # TODO - Need to implement this on the server side
-
-    def visualize(self):
-        """Visualize the token bucket over time."""
-        status = asyncio.run(self._get_status())
-        times, tokens = zip(*status["log"])
-        start_time = times[0]
-        times = [t - start_time for t in times]
-
-        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()
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod()
-    # bucket = TokenBucketClient(
-    #     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)
-    # print(throughput)
-    # bucket.visualize()

edsl/jobs/decorators.py

@@ -1,35 +0,0 @@
-from functools import wraps
-from threading import RLock
-import inspect
-
-
-def synchronized_class(wrapped_class):
-    """Class decorator that makes all methods thread-safe."""
-
-    # Add a lock to the class
-    setattr(wrapped_class, "_lock", RLock())
-
-    # Get all methods from the class
-    for name, method in inspect.getmembers(wrapped_class, inspect.isfunction):
-        # Skip magic methods except __getitem__, __setitem__, __delitem__
-        if name.startswith("__") and name not in [
-            "__getitem__",
-            "__setitem__",
-            "__delitem__",
-        ]:
-            continue
-
-        # Create synchronized version of the method
-        def create_synchronized_method(method):
-            @wraps(method)
-            def synchronized_method(*args, **kwargs):
-                instance = args[0]  # first arg is self
-                with instance._lock:
-                    return method(*args, **kwargs)
-
-            return synchronized_method
-
-        # Replace the original method with synchronized version
-        setattr(wrapped_class, name, create_synchronized_method(method))
-
-    return wrapped_class

edsl/jobs/jobs_status_enums.py

@@ -1,9 +0,0 @@
-from enum import Enum
-
-
-class JobsStatus(Enum):
-    QUEUED = "queued"
-    RUNNING = "running"
-    COMPLETED = "completed"
-    FAILED = "failed"
-    CANCELLED = "cancelled"