pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/exec/expr_eval/schedulers.py

@@ -4,9 +4,10 @@ import asyncio
 import datetime
 import inspect
 import logging
+import re
 import sys
 import time
-from typing import Awaitable, Collection, Optional
+from typing import Any, Awaitable, Collection
 
 from pixeltable import env, func
 from pixeltable.config import Config
@@ -34,7 +35,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,10 +77,12 @@ 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
 
@@ -96,12 +99,11 @@ 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)
             aws: list[Awaitable[None]] = []
-            completed_aw: Optional[asyncio.Task] = None
-            wait_for_reset: Optional[asyncio.Task] = None
+            completed_aw: asyncio.Task | None = None
+            wait_for_reset: asyncio.Task | None = None
             if limits_info is not None:
                 # limits_info's resource is depleted, wait for capacity to free up
 
@@ -115,21 +117,31 @@ class RateLimitsScheduler(Scheduler):
                 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()))
+                    wait_duration = (reset_at - now).total_seconds()
+                    wait_for_reset = asyncio.create_task(asyncio.sleep(wait_duration))
                     aws.append(wait_for_reset)
-                    _logger.debug(f'waiting for rate limit reset for {self.resource_pool}')
+                    _logger.debug(
+                        f'waiting {wait_duration:.2f}s for rate limit reset of '
+                        f'{self.resource_pool}:{limits_info.resource} (remaining={limits_info.remaining})'
+                    )
 
             if len(aws) > 0:
                 # we have something to wait for
+                report_ts = limits_info.recorded_at
                 done, pending = await asyncio.wait(aws, return_when=asyncio.FIRST_COMPLETED)
                 for task in pending:
                     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()
+                    _logger.debug(f'wait(): rate limit reset for {self.resource_pool}:{limits_info.resource}')
+                    last_report_ts = self.pool_info.resource_limits[limits_info.resource].recorded_at
+                    if report_ts == last_report_ts:
+                        # if we haven't seen a new report since we started waiting, force waiting for another rate limit
+                        # report before making any scheduling decisions
+                        # TODO: is it a good idea to discard the information we have?
+                        _logger.debug(f'resetting {self.resource_pool}: currently at {self.pool_info.debug_str()}')
+                        self.pool_info.reset()
                 # re-evaluate current capacity for current item
                 continue
 
@@ -155,18 +167,24 @@ 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]:
+    def _check_resource_limits(self, request_resources: dict[str, int]) -> env.RateLimitInfo | None:
         """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)
+        candidates: list[tuple[env.RateLimitInfo, float]] = []  # (info, relative remaining)
         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]
+            candidates.append((info, est_remaining / info.limit))
+        assert len(candidates) > 0
+        candidates.sort(key=lambda x: x[1])  # most depleted first
+        most_depleted = candidates[0]
+        _logger.debug(
+            f'check_resource_limits({request_resources}): '
+            f'most_depleted={most_depleted[0].resource}, rel_remaining={most_depleted[1]}'
+        )
+        # 0.05: leave some headroom, we don't have perfect information
+        if most_depleted[1] < 0.05:
+            return most_depleted[0]
+        return None
 
     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)
@@ -187,7 +205,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(
@@ -201,10 +220,14 @@ 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 hasattr(exc, 'response') and hasattr(exc.response, 'headers'):
+                _logger.debug(f'scheduler {self.resource_pool}: exception headers: {exc.response.headers}')
             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
+            self.pool_info.record_exc(exc)
+
             if num_retries < self.MAX_RETRIES:
                 retry_delay = self.pool_info.get_retry_delay(exc)
                 if retry_delay is not None:
@@ -213,7 +236,6 @@ class RateLimitsScheduler(Scheduler):
                     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
 
             # record the exception
             _, _, exc_tb = sys.exc_info()
@@ -248,10 +270,23 @@ 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
+    RATE_LIMIT_INDICATORS = ('rate limit', 'too many requests', '429', 'quota exceeded', 'throttled', 'rate exceeded')
+    RETRY_AFTER_PATTERNS = (
+        r'retry after (\d+(?:\.\d+)?)\s*seconds?',
+        r'try again in (\d+(?:\.\d+)?)\s*seconds?',
+        r'wait (\d+(?:\.\d+)?)\s*seconds?',
+        r'retry-after:\s*(\d+(?:\.\d+)?)',
+    )
+
+    # Exponential backoff defaults
+    BASE_RETRY_DELAY = 1.0  # in seconds
+    MAX_RETRY_DELAY = 60.0  # in seconds
+    RETRY_BACKOFF_MULTIPLIER = 2.0
 
     def __init__(self, resource_pool: str, dispatcher: Dispatcher):
         super().__init__(resource_pool, dispatcher)
@@ -260,6 +295,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 +314,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 +328,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 +378,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.debug(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_rate_limit_error, retry_after = self._is_rate_limit_error(exc)
+            if is_rate_limit_error 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 +400,124 @@ 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 _is_rate_limit_error(self, exc: Exception) -> tuple[bool, float | None]:
+        """Returns True if the exception indicates a rate limit error, and the retry delay in seconds."""
+        from http import HTTPStatus
+
+        # Check for HTTP status TOO_MANY_REQUESTS in various exception classes.
+        # We look for attributes that contain status codes, instead of checking the type of the exception,
+        # in order to handle a wider variety of exception classes.
+        is_rate_limit_error = False
+        retry_delay: float | None = None
+
+        # requests.HTTPError/httpx.HTTPStatusError
+        if (
+            hasattr(exc, 'response')
+            and hasattr(exc.response, 'status_code')
+            and exc.response.status_code == HTTPStatus.TOO_MANY_REQUESTS.value
+        ):
+            is_rate_limit_error = True
+            retry_delay = self._extract_retry_delay_from_headers(exc.response.headers)
+        elif (
+            # urllib.error.HTTPError
+            (hasattr(exc, 'code') and exc.code == HTTPStatus.TOO_MANY_REQUESTS.value)
+            # aiohttp.ClientResponseError
+            or (hasattr(exc, 'status') and exc.status == HTTPStatus.TOO_MANY_REQUESTS.value)
+        ) and hasattr(exc, 'headers'):
+            is_rate_limit_error = True
+            retry_delay = self._extract_retry_delay_from_headers(exc.headers)
+
+        if is_rate_limit_error:
+            return True, retry_delay
+
+        # Check common rate limit keywords in exception message
+        error_msg = str(exc).lower()
+        if any(indicator in error_msg for indicator in self.RATE_LIMIT_INDICATORS):
+            retry_delay = self._extract_retry_delay_from_message(error_msg)
+            return True, retry_delay
+
+        return False, None
+
+    def _extract_retry_delay_from_headers(self, headers: Any | None) -> float | None:
+        """Extract retry delay from HTTP headers."""
+        if headers is None:
+            return None
+
+        # convert headers to dict-like object for consistent access
+        header_dict: dict
+        if hasattr(headers, 'get'):
+            header_dict = headers
+        else:
+            # headers are a list of tuples or other format
+            try:
+                header_dict = dict(headers)
+            except (TypeError, ValueError):
+                return None
+        # normalize dict keys: lowercase and remove dashes
+        header_dict = {k.lower().replace('-', ''): v for k, v in header_dict.items()}
+
+        # check Retry-After header
+        retry_after = header_dict.get('retryafter')
+        if retry_after is not None:
+            try:
+                return float(retry_after)
+            except (ValueError, TypeError):
+                pass
+
+        # check X-RateLimit-Reset (Unix timestamp)
+        reset_time = header_dict.get('xratelimitreset')
+        if reset_time is not None:
+            try:
+                reset_timestamp = float(reset_time)
+                delay = max(0, reset_timestamp - time.time())
+                return delay
+            except (ValueError, TypeError):
+                pass
+
+        # check X-RateLimit-Reset-After (seconds from now)
+        reset_after = header_dict.get('xratelimitresetafter')
+        if reset_after is not None:
+            try:
+                return float(reset_after)
+            except (ValueError, TypeError):
+                pass
+
+        return None
+
+    def _extract_retry_delay_from_message(self, msg: str) -> float | None:
+        msg_lower = msg.lower()
+        for pattern in self.RETRY_AFTER_PATTERNS:
+            match = re.search(pattern, msg_lower)
+            if match is not None:
+                try:
+                    return float(match.group(1))
+                except (ValueError, TypeError):
+                    continue
+        return None
+
+    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:
+            delay = self.BASE_RETRY_DELAY * (self.RETRY_BACKOFF_MULTIPLIER**num_retries)
+            return max(min(delay, self.MAX_RETRY_DELAY), self.BASE_RETRY_DELAY)
+
 
 # all concrete Scheduler subclasses that implement matches()
 SCHEDULERS = [RateLimitsScheduler, RequestRateScheduler]

pixeltable/exec/globals.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import dataclasses
+
+from pixeltable.exprs import ArrayMd
+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
+
+    @classmethod
+    def from_dict(cls, d: dict) -> InlinedObjectMd:
+        if 'array_md' in d:
+            array_md = ArrayMd(**d['array_md'])
+            del d['array_md']
+            return cls(**d, array_md=array_md)
+        else:
+            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()
+        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]
@@ -42,7 +42,7 @@ class InMemoryDataNode(ExecNode):
         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