xtremeflow 0.1.0__py3-none-any.whl

xtremeflow/__init__.py

@@ -0,0 +1,5 @@
+'''
+XtremeFlow: A high-performance asynchronous task scheduler for LLM workloads.
+'''
+
+__version__ = '0.1.0'

xtremeflow/kvbatch.py

@@ -0,0 +1,59 @@
+'''Helper for KV cache-optimized async task batches.
+
+This module provides utilities for executing async tasks with a "first-wait,
+then-parallel" pattern optimized for KV cache utilization in LLM applications.
+
+Execution Pattern:
+
+    Input: [task1, task2, task3, ...]
+            ↓
+    ┌────────────────────────────────────┐
+    │ Phase 1: First Task                │
+    │ task1 runs to completion           │
+    │ (establishes KV cache)             │
+    └────────────────────────────────────┘
+            ↓
+    ┌────────────────────────────────────┐
+    │ Phase 2: Parallel Tasks            │
+    │ task2, task3, ... run concurrently │
+    │ (share the established cache)      │
+    └────────────────────────────────────┘
+            ↓
+    Output: [result1, result2, result3, ...]
+
+Use Case Example:
+
+    When scoring multiple resumes for the same job, each request shares the
+    job description prefix. The first request establishes a KV cache for the
+    job description. Subsequent requests can then run in parallel, leveraging
+    the cached computation for better performance.
+'''
+
+import asyncio
+from typing import Awaitable, Iterable, List, TypeVar
+
+T = TypeVar('T')
+
+
+async def _process_aws(*aws: Awaitable[T]) -> List[T]:
+    results = [await aws[0]] if aws else []
+    results += await asyncio.gather(*aws[1:])
+    return results
+
+
+def kv_batch(aws: Iterable[Awaitable[T]]) -> asyncio.Task[List[T]]:
+    '''Create a batch task with KV cache optimization.
+
+    Args:
+        aws: An iterable of awaitables to process.
+
+    Returns:
+        An asyncio.Task that completes with a list of results.
+
+    Example:
+        >>> task = kv_batch(
+        ...     llm_score(prompt) for prompt in same_job_with_different_resumes
+        ... )
+        >>> results = await task
+    '''
+    return asyncio.create_task(_process_aws(*aws))

xtremeflow/pipeline.py

@@ -0,0 +1,39 @@
+import asyncio
+from typing import Any, Callable, AsyncGenerator, AsyncIterable, Optional
+
+
+async def async_chunks(iterable: AsyncIterable, size: int):
+    it = aiter(iterable)
+    while True:
+        chunk = []
+        for _ in range(size):
+            try:
+                item = await anext(it)
+                chunk.append(item)
+            except StopAsyncIteration:
+                if chunk:
+                    yield chunk
+                return
+        yield chunk
+
+
+async def async_pipeline(
+    producer: Callable[[asyncio.Queue], Any], process_item: Optional[Callable[[Any], Any]] = None
+) -> AsyncGenerator[Any]:
+    queue = asyncio.Queue()
+
+    async def producer_wrapper():
+        await producer(queue)
+        await queue.put(None)
+
+    asyncio.create_task(producer_wrapper())
+
+    while True:
+        item = await queue.get()
+        if item is None:
+            break
+
+        try:
+            yield await process_item(item) if process_item else item
+        finally:
+            queue.task_done()

xtremeflow/scheduler/__init__.py

@@ -0,0 +1,5 @@
+from .base import TaskScheduler
+
+__all__ = [
+    'TaskScheduler',
+]

xtremeflow/scheduler/base.py

@@ -0,0 +1,35 @@
+import asyncio
+from typing import Any, Coroutine
+
+
+class TaskScheduler:
+    def __init__(self, max_concurrency: int):
+        if max_concurrency <= 0:
+            raise ValueError(f'max_concurrency must be positive, got {max_concurrency}')
+
+        self.semaphore = asyncio.Semaphore(max_concurrency)
+        self.active_tasks = 0
+        self.total_completed = 0
+        self.pending_tasks = set()
+
+    async def _execute_coro(self, coro: Coroutine, **kwargs) -> Any:
+        return await coro
+
+    def _task_done(self, task):
+        self.pending_tasks.discard(task)
+        self.active_tasks -= 1
+        self.total_completed += 1
+        self.semaphore.release()
+
+    async def start_task(self, coro: Coroutine, **kwargs) -> asyncio.Task:
+        await self.semaphore.acquire()
+        self.active_tasks += 1
+        task = asyncio.create_task(self._execute_coro(coro, **kwargs))
+        self.pending_tasks.add(task)
+        task.add_done_callback(self._task_done)
+        return task
+
+    async def wait_pending(self):
+        if self.pending_tasks:
+            await asyncio.gather(*self.pending_tasks)
+            self.pending_tasks.clear()

xtremeflow/scheduler/rate_limit.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from abc import ABC, abstractmethod
+from contextvars import ContextVar
+from dataclasses import dataclass
+from functools import wraps
+from typing import Any, Coroutine, Optional, Type, Union
+
+from .base import TaskScheduler
+
+logger = logging.getLogger(__name__)
+
+_current_ctx: ContextVar['Optional[ExecutionContext]'] = ContextVar('_current_ctx', default=None)
+
+
+@dataclass
+class ExecutionContext:
+    scheduler: RateLimitScheduler
+    extra: Optional[dict] = None
+
+
+class RetryException(Exception):
+    def __init__(self, message: str = '', retry_after: Optional[float] = None):
+        super().__init__(message)
+        self.retry_after = retry_after
+
+
+def auto_backoff(
+    retry_for: Union[Type[Exception], list[Type[Exception]], None] = None,
+    max_retries: int = 3,
+    base_retry_after: float = 2.0,
+    exponential: bool = True
+):
+    if retry_for is None:
+        retry_types = (RetryException,)
+    elif isinstance(retry_for, list):
+        retry_types = tuple(retry_for)
+    else:
+        retry_types = retry_for
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            last_exc = None
+            for attempt in range(max_retries + 1):
+                try:
+                    return await func(*args, **kwargs)
+                except retry_types as e:
+                    last_exc = e
+                    ctx = _current_ctx.get()
+
+                    if attempt < max_retries and ctx:
+                        header_wait = getattr(e, 'retry_after', None)
+                        if header_wait is not None and isinstance(header_wait, (int, float)):
+                            wait_sec = float(header_wait)
+                        else:
+                            wait_sec = base_retry_after * (2 ** attempt) if exponential else base_retry_after
+                        logger.warning(
+                            f'Retrying in {wait_sec:.1f}s '
+                            f'(attempt {attempt + 1}/{max_retries}): {e}'
+                        )
+                        ctx.scheduler.notify_rate_limit_exceeded(wait_sec)
+                        await asyncio.sleep(wait_sec)
+                        continue
+                    raise last_exc
+        return wrapper
+    return decorator
+
+
+def get_context() -> Optional[ExecutionContext]:
+    return _current_ctx.get()
+
+
+class RateLimitScheduler(TaskScheduler, ABC):
+    def __init__(self, max_concurrency: int, init_ratio: float = 0.0):
+        super().__init__(max_concurrency)
+        self._backoff_until = 0.0
+        self._initial_ratio = init_ratio
+
+    def notify_rate_limit_exceeded(self, retry_after: float):
+        self._backoff_until = max(self._backoff_until, time.monotonic() + retry_after)
+
+    def _get_backoff_wait(self) -> float:
+        return max(0.0, self._backoff_until - time.monotonic())
+
+    def _get_wait_time(self) -> float:
+        return self._get_backoff_wait()
+
+    @abstractmethod
+    def _consume_rate_quota(self):
+        '''Consume rate limit quota. Subclasses must implement this.'''
+
+    async def _wait_for_quota(self):
+        while True:
+            wait_time = self._get_wait_time()
+            if wait_time <= 0:
+                self._consume_rate_quota()
+                break
+            await asyncio.sleep(wait_time)
+
+    async def _execute_coro(self, coro: Coroutine, ctx_extra=None, **kwargs) -> Any:
+        ctx = ExecutionContext(scheduler=self, extra=ctx_extra)
+        token = _current_ctx.set(ctx)
+        try:
+            await self._wait_for_quota()
+            return await super()._execute_coro(coro, **kwargs)
+        finally:
+            _current_ctx.reset(token)

xtremeflow/scheduler/request.py

@@ -0,0 +1,49 @@
+import time
+from typing import Optional
+
+from .rate_limit import RateLimitScheduler
+
+
+class RequestRateScheduler(RateLimitScheduler):
+    def __init__(
+        self,
+        max_rps: Optional[int] = None,
+        max_rpm: Optional[int] = None,
+        *args,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self._max_rps = max_rps
+        self._max_rpm = max_rpm
+        self._rps_bucket = (max_rps or 0) * self._initial_ratio
+        self._rpm_bucket = (max_rpm or 0) * self._initial_ratio
+        self._last_req_update = time.monotonic()
+
+    def _get_wait_time(self) -> float:
+        now = time.monotonic()
+        delta = now - self._last_req_update
+        self._last_req_update = now
+
+        if self._max_rps:
+            self._rps_bucket = min(float(self._max_rps), self._rps_bucket + delta * self._max_rps)
+        if self._max_rpm:
+            self._rpm_bucket = min(float(self._max_rpm), self._rpm_bucket + delta * (self._max_rpm / 60.0))
+
+        waits = [super()._get_wait_time()]
+        if self._max_rps and self._rps_bucket < 1:
+            waits.append((1 - self._rps_bucket) / self._max_rps)
+        if self._max_rpm and self._rpm_bucket < 1:
+            waits.append((1 - self._rpm_bucket) / (self._max_rpm / 60.0))
+
+        return max(waits)
+
+    def _consume_rate_quota(self):
+        if self._max_rps:
+            self._rps_bucket -= 1
+        if self._max_rpm:
+            self._rpm_bucket -= 1
+
+    def reset_quota(self):
+        self._rps_bucket = (self._max_rps or 0) * self._initial_ratio
+        self._rpm_bucket = (self._max_rpm or 0) * self._initial_ratio
+        self._last_req_update = time.monotonic()

xtremeflow/scheduler/token.py

@@ -0,0 +1,79 @@
+import asyncio
+import time
+from typing import Optional, cast
+
+from .request import RequestRateScheduler
+from .rate_limit import get_context
+
+
+class TokenRateScheduler(RequestRateScheduler):
+    def __init__(
+        self,
+        max_tps: Optional[int] = None,
+        max_tpm: Optional[int] = None,
+        *args,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self._max_tps = max_tps
+        self._max_tpm = max_tpm
+        self._tps_bucket = (self._max_tps or 0) * self._initial_ratio
+        self._tpm_bucket = (self._max_tpm or 0) * self._initial_ratio
+        self._last_token_update = time.monotonic()
+
+    async def start_task(self, coro, estimated_tokens: int, **kwargs) -> asyncio.Task:
+        return await super().start_task(
+            coro, ctx_extra={'estimated_tokens': estimated_tokens},
+            **kwargs)
+
+    def _get_wait_time(self) -> float:
+        now = time.monotonic()
+        delta = now - self._last_token_update
+        self._last_token_update = now
+
+        if self._max_tps:
+            self._tps_bucket = min(float(self._max_tps), self._tps_bucket + delta * self._max_tps)
+        if self._max_tpm:
+            self._tpm_bucket = min(float(self._max_tpm), self._tpm_bucket + delta * (self._max_tpm / 60.0))
+
+        waits = [super()._get_wait_time()]
+
+        ctx = get_context()
+        tokens = ctx.extra.get('estimated_tokens', 0)
+        if tokens > 0:
+            if self._max_tps and self._tps_bucket < tokens:
+                waits.append((tokens - self._tps_bucket) / self._max_tps)
+            if self._max_tpm and self._tpm_bucket < tokens:
+                waits.append((tokens - self._tpm_bucket) / (self._max_tpm / 60.0))
+        return max(waits)
+
+    def _consume_rate_quota(self):
+        super()._consume_rate_quota()
+        ctx = get_context()
+        tokens = ctx.extra.get('estimated_tokens', 0)
+        if tokens > 0:
+            if self._max_tps:
+                self._tps_bucket -= tokens
+            if self._max_tpm:
+                self._tpm_bucket -= tokens
+
+    def _apply_correction(self, actual: int):
+        ctx = get_context()
+        estimated = ctx.extra.get('estimated_tokens', 0)
+        diff = estimated - actual
+        if diff == 0:
+            return
+        if self._max_tps:
+            self._tps_bucket = min(float(self._max_tps), self._tps_bucket + diff)
+        if self._max_tpm:
+            self._tpm_bucket = min(float(self._max_tpm), self._tpm_bucket + diff)
+
+    def reset_quota(self):
+        self._tps_bucket = (self._max_tps or 0) * self._initial_ratio
+        self._tpm_bucket = (self._max_tpm or 0) * self._initial_ratio
+        self._last_token_update = time.monotonic()
+
+
+def report_token_usage(actual: int):
+    ctx = get_context()
+    cast(TokenRateScheduler, ctx.scheduler)._apply_correction(actual)

xtremeflow-0.1.0.dist-info/METADATA

@@ -0,0 +1,139 @@
+Metadata-Version: 2.4
+Name: xtremeflow
+Version: 0.1.0
+Summary: XtremeFlow: A high-performance Python asynchronous task scheduler engineered to push LLM workloads to their absolute physical limits
+Author-email: Flow Jiang <flowjzh@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/flowjzh/xtremeflow
+Project-URL: Repository, https://github.com/flowjzh/xtremeflow.git
+Project-URL: Issues, https://github.com/flowjzh/xtremeflow/issues
+Keywords: async,scheduler,rate-limiting,llm,asyncio,concurrency,backpressure
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Operating System :: OS Independent
+Classifier: Typing :: Typed
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: dev
+Requires-Dist: pytest>=8.4.2; extra == "dev"
+Requires-Dist: pytest-asyncio>=1.2.0; extra == "dev"
+Dynamic: license-file
+
+# XtremeFlow
+
+> **"Exhaust rate limits, not patience. Squeezing maximum throughput from every second."**
+
+### 🦅 About
+
+**XtremeFlow** is a high-performance asynchronous task scheduler engineered to push **Large Language Model (LLM)** workloads to their absolute physical limits.
+
+**The Problem:**
+LLM providers throttle your velocity through a combination of **Concurrency**, **RPS**/**RPM** or **TPS**/**TPM**. Most schedulers are defensive—they wait too long, leave gaps in your schedule, and waste capacity. In high-volume production, idle time is a lost resource.
+
+**The XtremeFlow Philosophy:**
+Stop being polite with your rate limits. **XtremeFlow is offensive.** It is designed to saturate your provider's capacity with surgical precision. Using a unique **Backpressure Reflex**, it maintains peak velocity until the very moment a limit is hit, executes a synchronized global cool-down, and resumes at full speed the millisecond the provider allows.
+
+> ⚠️ **Limitation:** XtremeFlow is currently optimized for **single-process** `asyncio` applications. It manages state in-memory and does not support distributed rate limiting (e.g., Redis-based) out of the box.
+
+### ⚡ Key Features
+
+* **Aggressive Saturation**: Engineered to fill every available millisecond of your allowed rate, ensuring zero wasted throughput.
+* **Backpressure Reflex**: Automatically detects 429 triggers and orchestrates a global **Exponential Backoff** across all workers to stay in perfect sync with provider resets.
+* **Dynamic Calibration**: Supports post-request reporting of *actual* usage to instantly "refund" over-estimated capacity back to the scheduler.
+* **Async-Native**: Built on `asyncio` for low-latency scheduling where every microsecond counts.
+* **KV Cache Optimization**: Provides utilities to maximize KV cache utilization across parallel LLM requests, dramatically reducing token consumption and improving throughput.
+* **Async Pipeline**: Producer-consumer pipeline for streaming workloads with automatic backpressure handling.
+
+### 🚀 Quick Start
+
+```python
+import asyncio
+from openai import RateLimitError
+from xtremeflow.scheduler.rate_limit import auto_backoff, report_token_usage
+from xtremeflow.scheduler.token import TokenRateScheduler
+
+# Initialize: 10 concurrent slots, 60 RPM, 50k TPM
+scheduler = TokenRateScheduler(
+    max_concurrency=10,
+    rpm=60,
+    tpm=50000
+)
+
+@auto_backoff(retry_for=RateLimitError, base_retry_after=2.0)
+async def call_llm_api(prompt: str):
+    """
+    Wraps LLM call with Backpressure Reflex.
+    Global synchronization ensures you don't keep hitting the wall during cooldown.
+    """
+    print(f"Executing task: {prompt}")
+    
+    # Simulated API call
+    await asyncio.sleep(1)
+    
+    # Calibration: Refund unused quota to the scheduler
+    report_token_usage(actual_tokens=450)
+    
+    return "success"
+
+async def main():
+    tasks = []
+    for i in range(10):
+        # Dispatch with an estimated cost to saturate the current limit
+        t = await scheduler.start_task(
+            call_llm_api(f"Task {i}"), 
+            estimated_tokens=500
+        )
+        tasks.append(t)
+    
+    results = await asyncio.gather(*tasks)
+    print(f"XtremeFlow: Successfully processed {len(results)} tasks at peak throughput.")
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+### 🔥 Performance Tools
+
+Beyond rate limiting, XtremeFlow provides utilities to maximize token efficiency and throughput.
+
+**KV Cache Optimization** (`kv_batch`)
+```python
+from xtremeflow.kvbatch import kv_batch
+
+# First request establishes KV cache, rest run in parallel
+task = kv_batch(
+    llm_score(prompt) for prompt in same_job_with_different_resumes
+)
+results = await task
+```
+Reduces token consumption by 40-60% for batched requests with shared prefixes.
+
+**Async Pipeline** (`async_pipeline`)
+```python
+from xtremeflow.pipeline import async_pipeline
+
+# Producer: scheduler-controlled, exhausts this tier's rate limit
+async def producer(queue: asyncio.Queue):
+    async for item in source:
+        task = await scheduler.start_task(llm_api(item), estimate_tokens)
+        await queue.put(task)
+
+# Processor: slower sequential processing, yields to next tier
+async def process_item(item):
+    result = await item
+    return await db_write(result)  # Different rate limit tier
+
+async for result in async_pipeline(producer, process_item):
+    yield result  # Can chain to another tier
+```
+Decouples rate limit tiers—exhausting each tier's limit frees up quota for other tasks immediately, maximizing overall system throughput.

xtremeflow-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+xtremeflow/__init__.py,sha256=nSaIwlYraiF6jrhcHU5rDDa7c_PXYk-TGEJ85NRK1xM,109
+xtremeflow/kvbatch.py,sha256=F_fhRn4p3t4yXDbNAu-_GlnbjwqZCC6LFgZdygTn2Rw,2173
+xtremeflow/pipeline.py,sha256=RIblYq9cmjbUPnq3KbxliMx34Qj6BAvrdRV2H-tApWs,1015
+xtremeflow/scheduler/__init__.py,sha256=zFFT4IWQfZtnt1bGwbAITv4W8jMK96W3UrSk7-EV10U,68
+xtremeflow/scheduler/base.py,sha256=nfo4B90VcEem55AyF21kXQA8TUgiUtIOKoUXTPSpSiA,1148
+xtremeflow/scheduler/rate_limit.py,sha256=f9L_eCXC2Ess9vl8n90PY5_i4yu7zjiNOR9trjhCov4,3711
+xtremeflow/scheduler/request.py,sha256=OHR3WcsxgTTsvDcB8rkxeldfwUK57UZa3f5OtTy-T1U,1662
+xtremeflow/scheduler/token.py,sha256=4qWvOZvWNy3Nd4tiEuCCy_9kigK1rVvvdls0EQKetI8,2822
+xtremeflow-0.1.0.dist-info/licenses/LICENSE,sha256=OhxWdg6CEDY67lRrpbPvRuKrxRpcZEL4S84Ip25MFGk,1067
+xtremeflow-0.1.0.dist-info/METADATA,sha256=tMngmtrwxoafokwDSyBpuGmWNb82-m3kNktd2WOEQ90,5969
+xtremeflow-0.1.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+xtremeflow-0.1.0.dist-info/top_level.txt,sha256=eBfGcIsAFKFX4AKTq_ELDu8YoZhhCJWv3ijTjFBiVWs,11
+xtremeflow-0.1.0.dist-info/RECORD,,

xtremeflow-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

xtremeflow-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Flow Jiang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

xtremeflow-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+xtremeflow