edsl 0.1.39__py3-none-any.whl → 0.1.39.dev1__py3-none-any.whl

edsl/jobs/async_interview_runner.py

@@ -1,138 +0,0 @@
-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/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/check_survey_scenario_compatibility.py

@@ -1,85 +0,0 @@
-import warnings
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from edsl.surveys.Survey import Survey
-    from edsl.scenarios.ScenarioList import ScenarioList
-
-
-class CheckSurveyScenarioCompatibility:
-
-    def __init__(self, survey: "Survey", scenarios: "ScenarioList"):
-        self.survey = survey
-        self.scenarios = scenarios
-
-    def check(self, strict: bool = False, warn: bool = False) -> None:
-        """Check if the parameters in the survey and scenarios are consistent.
-
-        >>> from edsl.jobs.Jobs import Jobs
-        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
-        >>> from edsl.surveys.Survey import Survey
-        >>> from edsl.scenarios.Scenario import Scenario
-        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
-        >>> j = Jobs(survey = Survey(questions=[q]))
-        >>> cs = CheckSurveyScenarioCompatibility(j.survey, j.scenarios)
-        >>> with warnings.catch_warnings(record=True) as w:
-        ...     cs.check(warn = True)
-        ...     assert len(w) == 1
-        ...     assert issubclass(w[-1].category, UserWarning)
-        ...     assert "The following parameters are in the survey but not in the scenarios" in str(w[-1].message)
-
-        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
-        >>> s = Scenario({'plop': "A", 'poo': "B"})
-        >>> j = Jobs(survey = Survey(questions=[q])).by(s)
-        >>> cs = CheckSurveyScenarioCompatibility(j.survey, j.scenarios)
-        >>> cs.check(strict = True)
-        Traceback (most recent call last):
-        ...
-        ValueError: The following parameters are in the scenarios but not in the survey: {'plop'}
-
-        >>> q = QuestionFreeText(question_text = "Hello", question_name = "ugly_question")
-        >>> s = Scenario({'ugly_question': "B"})
-        >>> from edsl.scenarios.ScenarioList import ScenarioList
-        >>> cs = CheckSurveyScenarioCompatibility(Survey(questions=[q]), ScenarioList([s]))
-        >>> cs.check()
-        Traceback (most recent call last):
-        ...
-        ValueError: The following names are in both the survey question_names and the scenario keys: {'ugly_question'}. This will create issues.
-        """
-        survey_parameters: set = self.survey.parameters
-        scenario_parameters: set = self.scenarios.parameters
-
-        msg0, msg1, msg2 = None, None, None
-
-        # look for key issues
-        if intersection := set(self.scenarios.parameters) & set(
-            self.survey.question_names
-        ):
-            msg0 = f"The following names are in both the survey question_names and the scenario keys: {intersection}. This will create issues."
-
-            raise ValueError(msg0)
-
-        if in_survey_but_not_in_scenarios := survey_parameters - scenario_parameters:
-            msg1 = f"The following parameters are in the survey but not in the scenarios: {in_survey_but_not_in_scenarios}"
-        if in_scenarios_but_not_in_survey := scenario_parameters - survey_parameters:
-            msg2 = f"The following parameters are in the scenarios but not in the survey: {in_scenarios_but_not_in_survey}"
-
-        if msg1 or msg2:
-            message = "\n".join(filter(None, [msg1, msg2]))
-            if strict:
-                raise ValueError(message)
-            else:
-                if warn:
-                    warnings.warn(message)
-
-        if self.scenarios.has_jinja_braces:
-            warnings.warn(
-                "The scenarios have Jinja braces ({{ and }}). Converting to '<<' and '>>'. If you want a different conversion, use the convert_jinja_braces method first to modify the scenario."
-            )
-            self.scenarios = self.scenarios._convert_jinja_braces()
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod()