pixeltable 0.3.14__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/exec/expr_eval/schedulers.py

@@ -4,17 +4,21 @@ import asyncio
 import datetime
 import inspect
 import logging
+import math
 import sys
 import time
-from typing import Awaitable, Collection, Optional
+from typing import Awaitable, Collection
 
 from pixeltable import env, func
 from pixeltable.config import Config
+from pixeltable.utils.http import exponential_backoff, is_retriable_error
 
 from .globals import Dispatcher, ExecCtx, FnCallArgs, Scheduler
 
 _logger = logging.getLogger('pixeltable')
 
+__all__ = ['RateLimitsScheduler', 'RequestRateScheduler']
+
 
 class RateLimitsScheduler(Scheduler):
     """
@@ -34,7 +38,7 @@ class RateLimitsScheduler(Scheduler):
     get_request_resources_param_names: list[str]  # names of parameters of RateLimitsInfo.get_request_resources()
 
     # scheduling-related state
-    pool_info: Optional[env.RateLimitsInfo]
+    pool_info: env.RateLimitsInfo | None
     est_usage: dict[str, int]  # value per resource; accumulated estimates since the last util. report
 
     num_in_flight: int  # unfinished tasks
@@ -76,14 +80,15 @@ class RateLimitsScheduler(Scheduler):
         self.est_usage = dict.fromkeys(self._resources, 0)
 
     async def _main_loop(self) -> None:
-        item: Optional[RateLimitsScheduler.QueueItem] = None
+        item: RateLimitsScheduler.QueueItem | None = None
         while True:
             if item is None:
                 item = await self.queue.get()
+                assert isinstance(item.request.fn_call.fn, func.CallableFunction)
+                assert '_runtime_ctx' in item.request.fn_call.fn.signature.system_parameters
                 if item.num_retries > 0:
                     self.total_retried += 1
 
-            now = datetime.datetime.now(tz=datetime.timezone.utc)
             if self.pool_info is None or not self.pool_info.is_initialized():
                 # wait for a single request to get rate limits
                 _logger.debug(f'initializing rate limits for {self.resource_pool}')
@@ -96,14 +101,13 @@ class RateLimitsScheduler(Scheduler):
                 continue
 
             # check rate limits
-            _logger.debug(f'checking rate limits for {self.resource_pool}')
             request_resources = self._get_request_resources(item.request)
-            limits_info = self._check_resource_limits(request_resources)
+            resource_delay = self._resource_delay(request_resources)
             aws: list[Awaitable[None]] = []
-            completed_aw: Optional[asyncio.Task] = None
-            wait_for_reset: Optional[asyncio.Task] = None
-            if limits_info is not None:
-                # limits_info's resource is depleted, wait for capacity to free up
+            completed_aw: asyncio.Task | None = None
+            wait_for_reset: asyncio.Task | None = None
+            if resource_delay > 0:
+                # Some resource or resources are nearing depletion
 
                 if self.num_in_flight > 0:
                     # a completed request can free up capacity
@@ -112,12 +116,10 @@ class RateLimitsScheduler(Scheduler):
                     aws.append(completed_aw)
                     _logger.debug(f'waiting for completed request for {self.resource_pool}')
 
-                reset_at = limits_info.reset_at
-                if reset_at > now:
-                    # we're waiting for the rate limit to reset
-                    wait_for_reset = asyncio.create_task(asyncio.sleep((reset_at - now).total_seconds()))
-                    aws.append(wait_for_reset)
-                    _logger.debug(f'waiting for rate limit reset for {self.resource_pool}')
+                # Schedule a sleep until sufficient resources are available
+                wait_for_reset = asyncio.create_task(asyncio.sleep(resource_delay))
+                aws.append(wait_for_reset)
+                _logger.debug(f'waiting {resource_delay:.1f}s for resource availability')
 
             if len(aws) > 0:
                 # we have something to wait for
@@ -126,16 +128,12 @@ class RateLimitsScheduler(Scheduler):
                     task.cancel()
                 if completed_aw in done:
                     _logger.debug(f'wait(): completed request for {self.resource_pool}')
-                if wait_for_reset in done:
-                    _logger.debug(f'wait(): rate limit reset for {self.resource_pool}')
-                    # force waiting for another rate limit report before making any scheduling decisions
-                    self.pool_info.reset()
                 # re-evaluate current capacity for current item
                 continue
 
             # we have a new in-flight request
             for resource, val in request_resources.items():
-                self.est_usage[resource] += val
+                self.est_usage[resource] = self.est_usage.get(resource, 0) + val
             _logger.debug(f'creating task for {self.resource_pool}')
             self.num_in_flight += 1
             task = asyncio.create_task(self._exec(item.request, item.exec_ctx, item.num_retries, is_task=True))
@@ -155,25 +153,30 @@ class RateLimitsScheduler(Scheduler):
             constant_kwargs, batch_kwargs = request.pxt_fn.create_batch_kwargs(batch_kwargs)
             return self.pool_info.get_request_resources(**constant_kwargs, **batch_kwargs)
 
-    def _check_resource_limits(self, request_resources: dict[str, int]) -> Optional[env.RateLimitInfo]:
-        """Returns the most depleted resource, relative to its limit, or None if all resources are within limits"""
-        candidates: list[tuple[env.RateLimitInfo, float]] = []  # (info, relative usage)
+    def _resource_delay(self, request_resources: dict[str, int]) -> float:
+        """For the provided resources and usage, attempts to estimate the time to wait until sufficient resources are
+        available."""
+        highest_wait = 0.0
+        highest_wait_resource = None
         for resource, usage in request_resources.items():
-            # 0.05: leave some headroom, we don't have perfect information
             info = self.pool_info.resource_limits[resource]
-            est_remaining = info.remaining - self.est_usage[resource] - usage
-            if est_remaining < 0.05 * info.limit:
-                candidates.append((info, est_remaining / info.limit))
-        if len(candidates) == 0:
-            return None
-        return min(candidates, key=lambda x: x[1])[0]
+            # Note: usage and est_usage are estimated costs of requests, and it may be way off (for example, if max
+            # tokens is unspecified for an openAI request).
+            time_until = info.estimated_resource_refill_delay(
+                math.ceil(info.limit * env.TARGET_RATE_LIMIT_RESOURCE_FRACT + usage + self.est_usage.get(resource, 0))
+            )
+            if time_until is not None and highest_wait < time_until:
+                highest_wait = time_until
+                highest_wait_resource = resource
+        _logger.debug(f'Determined wait time of {highest_wait:.1f}s for resource {highest_wait_resource}')
+        return highest_wait
 
     async def _exec(self, request: FnCallArgs, exec_ctx: ExecCtx, num_retries: int, is_task: bool) -> None:
         assert all(not row.has_val[request.fn_call.slot_idx] for row in request.rows)
         assert all(not row.has_exc(request.fn_call.slot_idx) for row in request.rows)
 
+        start_ts = datetime.datetime.now(tz=datetime.timezone.utc)
         try:
-            start_ts = datetime.datetime.now(tz=datetime.timezone.utc)
             pxt_fn = request.fn_call.fn
             assert isinstance(pxt_fn, func.CallableFunction)
             _logger.debug(
@@ -187,7 +190,8 @@ class RateLimitsScheduler(Scheduler):
                 for row, result in zip(request.rows, batch_result):
                     row[request.fn_call.slot_idx] = result
             else:
-                result = await pxt_fn.aexec(*request.args, **request.kwargs)
+                request_kwargs = {**request.kwargs, '_runtime_ctx': env.RuntimeCtx(is_retry=num_retries > 0)}
+                result = await pxt_fn.aexec(*request.args, **request_kwargs)
                 request.row[request.fn_call.slot_idx] = result
             end_ts = datetime.datetime.now(tz=datetime.timezone.utc)
             _logger.debug(
@@ -200,20 +204,32 @@ class RateLimitsScheduler(Scheduler):
 
             self.dispatcher.dispatch(request.rows, exec_ctx)
         except Exception as exc:
-            _logger.debug(f'scheduler {self.resource_pool}: exception in slot {request.fn_call.slot_idx}: {exc}')
-            if self.pool_info is None:
-                # our pool info should be available at this point
-                self._set_pool_info()
-            assert self.pool_info is not None
-            if num_retries < self.MAX_RETRIES:
-                retry_delay = self.pool_info.get_retry_delay(exc)
-                if retry_delay is not None:
-                    self.total_retried += 1
-                    _logger.debug(f'scheduler {self.resource_pool}: retrying in {retry_delay} seconds')
-                    await asyncio.sleep(retry_delay)
-                    self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx))
-                    return
-            # TODO: update resource limits reported in exc.response.headers, if present
+            _logger.exception(f'scheduler {self.resource_pool}: exception in slot {request.fn_call.slot_idx}: {exc}')
+            if hasattr(exc, 'response') and hasattr(exc.response, 'headers'):
+                _logger.debug(f'scheduler {self.resource_pool}: exception headers: {exc.response.headers}')
+
+            # If pool info is available, attempt to retry based on the resource information
+            # Pool info may not be available yet if the exception occurred before the UDF set it
+            if self.pool_info is not None:
+                self.pool_info.record_exc(start_ts, exc)
+
+                if num_retries < self.MAX_RETRIES:
+                    retry_delay = self.pool_info.get_retry_delay(exc, num_retries)
+                    if retry_delay is None:
+                        # The resource pool did not recognize it as a retriable error. Try our generic best-effort logic
+                        # before giving up.
+                        is_retriable, retry_delay = is_retriable_error(exc)
+                        if is_retriable:
+                            retry_delay = retry_delay or exponential_backoff(num_retries)
+                    if retry_delay is not None:
+                        self.total_retried += 1
+                        _logger.debug(
+                            f'scheduler {self.resource_pool}: sleeping {retry_delay:.2f}s before retrying'
+                            f' attempt {num_retries} based on the information in the error'
+                        )
+                        await asyncio.sleep(retry_delay)
+                        self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx))
+                        return
 
             # record the exception
             _, _, exc_tb = sys.exc_info()
@@ -248,11 +264,16 @@ class RequestRateScheduler(Scheduler):
     num_in_flight: int
     total_requests: int
     total_retried: int
+    total_errors: int
 
     TIME_FORMAT = '%H:%M.%S %f'
-    MAX_RETRIES = 10
+    MAX_RETRIES = 3
     DEFAULT_RATE_LIMIT = 600  # requests per minute
 
+    # Exponential backoff defaults
+    BASE_RETRY_DELAY = 1.0  # in seconds
+    MAX_RETRY_DELAY = 60.0  # in seconds
+
     def __init__(self, resource_pool: str, dispatcher: Dispatcher):
         super().__init__(resource_pool, dispatcher)
         loop_task = asyncio.create_task(self._main_loop())
@@ -260,6 +281,7 @@ class RequestRateScheduler(Scheduler):
         self.num_in_flight = 0
         self.total_requests = 0
         self.total_retried = 0
+        self.total_errors = 0
 
         # try to get the rate limit from the config
         elems = resource_pool.split(':')
@@ -278,6 +300,7 @@ class RequestRateScheduler(Scheduler):
             key = model
         requests_per_min = Config.get().get_int_value(key, section=section)
         requests_per_min = requests_per_min or self.DEFAULT_RATE_LIMIT
+        _logger.debug(f'rate limit for {self.resource_pool}: {requests_per_min} RPM')
         self.secs_per_request = 1 / (requests_per_min / 60)
 
     @classmethod
@@ -291,8 +314,12 @@ class RequestRateScheduler(Scheduler):
             if item.num_retries > 0:
                 self.total_retried += 1
             now = time.monotonic()
+            wait_duration = 0.0
+            if item.retry_after is not None:
+                wait_duration = item.retry_after - now
             if now - last_request_ts < self.secs_per_request:
-                wait_duration = self.secs_per_request - (now - last_request_ts)
+                wait_duration = max(wait_duration, self.secs_per_request - (now - last_request_ts))
+            if wait_duration > 0:
                 _logger.debug(f'waiting for {wait_duration} for {self.resource_pool}')
                 await asyncio.sleep(wait_duration)
 
@@ -337,15 +364,21 @@ class RequestRateScheduler(Scheduler):
             self.dispatcher.dispatch(request.rows, exec_ctx)
 
         except Exception as exc:
-            # TODO: which exception can be retried?
-            _logger.debug(f'exception for {self.resource_pool}: {exc}')
-            status = getattr(exc, 'status', None)
-            _logger.debug(f'type={type(exc)} has_status={hasattr(exc, "status")} status={status}')
-            if num_retries < self.MAX_RETRIES:
-                self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx))
+            _logger.exception(f'exception for {self.resource_pool}: type={type(exc)}\n{exc}')
+            if hasattr(exc, 'response') and hasattr(exc.response, 'headers'):
+                _logger.debug(f'scheduler {self.resource_pool}: exception headers: {exc.response.headers}')
+            is_retriable, retry_after = is_retriable_error(exc)
+            if is_retriable and num_retries < self.MAX_RETRIES:
+                retry_delay = self._compute_retry_delay(num_retries, retry_after)
+                _logger.debug(f'scheduler {self.resource_pool}: retrying after {retry_delay}')
+                now = time.monotonic()
+                # put the request back in the queue right away, which prevents new requests from being generated until
+                # this one succeeds or exceeds its retry limit
+                self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx, retry_after=now + retry_delay))
                 return
 
             # record the exception
+            self.total_errors += 1
             _, _, exc_tb = sys.exc_info()
             for row in request.rows:
                 row.set_exc(request.fn_call.slot_idx, exc)
@@ -353,11 +386,28 @@ class RequestRateScheduler(Scheduler):
         finally:
             _logger.debug(
                 f'Scheduler stats: #in-flight={self.num_in_flight} #requests={self.total_requests}, '
-                f'#retried={self.total_retried}'
+                f'#retried={self.total_retried} #errors={self.total_errors}'
             )
             if is_task:
                 self.num_in_flight -= 1
 
+    def _compute_retry_delay(self, num_retries: int, retry_after: float | None = None) -> float:
+        """
+        Calculate exponential backoff delay for rate limit errors.
+
+        Args:
+            retry_count: Number of retries attempted (0-based)
+            retry_after: Suggested delay from Retry-After header
+
+        Returns:
+            Delay in seconds
+        """
+        if retry_after is not None and retry_after > 0:
+            # Use server-suggested delay, but cap it at max_delay
+            return max(min(retry_after, self.MAX_RETRY_DELAY), self.BASE_RETRY_DELAY)
+        else:
+            return exponential_backoff(num_retries, max_delay=self.MAX_RETRY_DELAY)
+
 
 # all concrete Scheduler subclasses that implement matches()
 SCHEDULERS = [RateLimitsScheduler, RequestRateScheduler]

pixeltable/exec/globals.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import dataclasses
+
+from pixeltable.exprs import ArrayMd, BinaryMd
+from pixeltable.utils.misc import non_none_dict_factory
+
+INLINED_OBJECT_MD_KEY = '__pxtinlinedobjmd__'
+
+
+@dataclasses.dataclass
+class InlinedObjectMd:
+    type: str  # corresponds to ts.ColumnType.Type
+    url_idx: int
+    img_start: int | None = None
+    img_end: int | None = None
+    array_md: ArrayMd | None = None
+    binary_md: BinaryMd | None = None
+
+    @classmethod
+    def from_dict(cls, d: dict) -> InlinedObjectMd:
+        d = d.copy()
+        if 'array_md' in d:
+            d['array_md'] = ArrayMd(**d['array_md'])
+        if 'binary_md' in d:
+            d['binary_md'] = BinaryMd(**d['binary_md'])
+        return cls(**d)
+
+    def as_dict(self) -> dict:
+        result = dataclasses.asdict(self, dict_factory=non_none_dict_factory)
+        if self.array_md is not None:
+            result['array_md'] = self.array_md.as_dict()
+        if self.binary_md is not None:
+            result['binary_md'] = dataclasses.asdict(self.binary_md)
+        return result

pixeltable/exec/in_memory_data_node.py

@@ -1,8 +1,8 @@
 import logging
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncIterator
 
 from pixeltable import catalog, exprs
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.local_store import TempStore
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -23,7 +23,7 @@ class InMemoryDataNode(ExecNode):
 
     input_rows: list[dict[str, Any]]
     start_row_id: int
-    output_rows: Optional[DataRowBatch]
+    output_batch: DataRowBatch | None
 
     # output_exprs is declared in the superclass, but we redeclare it here with a more specific type
     output_exprs: list[exprs.ColumnRef]
@@ -38,11 +38,11 @@ class InMemoryDataNode(ExecNode):
         # we materialize the input slots
         output_exprs = list(row_builder.input_exprs)
         super().__init__(row_builder, output_exprs, [], None)
-        assert tbl.get().is_insertable()
+        assert tbl.get().is_insertable
         self.tbl = tbl
         self.input_rows = rows
         self.start_row_id = start_row_id
-        self.output_rows = None
+        self.output_batch = None
 
     def _open(self) -> None:
         """Create row batch and populate with self.input_rows"""
@@ -56,22 +56,21 @@ class InMemoryDataNode(ExecNode):
         }
         output_slot_idxs = {e.slot_idx for e in self.output_exprs}
 
-        self.output_rows = DataRowBatch(self.tbl, self.row_builder, len(self.input_rows))
-        for row_idx, input_row in enumerate(self.input_rows):
+        self.output_batch = DataRowBatch(self.row_builder)
+        for input_row in self.input_rows:
+            output_row = self.row_builder.make_row()
             # populate the output row with the values provided in the input row
             input_slot_idxs: set[int] = set()
             for col_name, val in input_row.items():
                 col_info = user_cols_by_name.get(col_name)
                 assert col_info is not None
-
-                if col_info.col.col_type.is_image_type() and isinstance(val, bytes):
-                    # this is a literal image, ie, a sequence of bytes; we save this as a media file and store the path
-                    path = str(MediaStore.prepare_media_path(self.tbl.id, col_info.col.id, self.tbl.get().version))
-                    with open(path, 'wb') as fp:
-                        fp.write(val)
-                    self.output_rows[row_idx][col_info.slot_idx] = path
+                col = col_info.col
+                if col.col_type.is_image_type() and isinstance(val, bytes):
+                    # this is a literal media file, ie, a sequence of bytes; save it as a binary file and store the path
+                    filepath, _ = TempStore.save_media_object(val, col, format=None)
+                    output_row[col_info.slot_idx] = str(filepath)
                 else:
-                    self.output_rows[row_idx][col_info.slot_idx] = val
+                    output_row[col_info.slot_idx] = val
 
                 input_slot_idxs.add(col_info.slot_idx)
 
@@ -80,10 +79,11 @@ class InMemoryDataNode(ExecNode):
             for slot_idx in missing_slot_idxs:
                 col_info = output_cols_by_idx.get(slot_idx)
                 assert col_info is not None
-                self.output_rows[row_idx][col_info.slot_idx] = None
+                output_row[col_info.slot_idx] = None
+            self.output_batch.add_row(output_row)
 
-        self.ctx.num_rows = len(self.output_rows)
+        self.ctx.num_rows = len(self.output_batch)
 
     async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
-        _logger.debug(f'InMemoryDataNode: created row batch with {len(self.output_rows)} output_rows')
-        yield self.output_rows
+        _logger.debug(f'InMemoryDataNode: created row batch with {len(self.output_batch)} rows')
+        yield self.output_batch