pixeltable 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

pixeltable/exec/expr_eval/expr_eval_node.py

@@ -4,24 +4,23 @@ import asyncio
 import logging
 import traceback
 from types import TracebackType
-from typing import Iterable, AsyncIterator, Optional, Union
+from typing import AsyncIterator, Iterable, Optional, Union
 
 import numpy as np
 
 import pixeltable.exceptions as excs
-from pixeltable import exprs
-from pixeltable import func
+from pixeltable import exprs, func
+
+from ..data_row_batch import DataRowBatch
+from ..exec_node import ExecNode
 from .evaluators import DefaultExprEvaluator, FnCallEvaluator
 from .globals import Evaluator, Scheduler
 from .row_buffer import RowBuffer
 from .schedulers import SCHEDULERS
-from ..data_row_batch import DataRowBatch
-from ..exec_node import ExecNode
 
 _logger = logging.getLogger('pixeltable')
 
 
-
 class ExprEvalNode(ExecNode):
     """
     Expression evaluation
@@ -35,10 +34,13 @@ class ExprEvalNode(ExecNode):
     TODO:
     - Literal handling: currently, Literal values are copied into slots via the normal evaluation mechanism, which is
       needless overhead; instead: pre-populate Literal slots in _init_row()
+    - dynamically determine MAX_BUFFERED_ROWS, based on the avg memory consumption of a row and our configured memory
+      limit
     - local model inference on gpu: currently, no attempt is made to ensure that models can fit onto the gpu
       simultaneously, which will cause errors; instead, the execution should be divided into sequential phases, each
       of which only contains a subset of the models which is known to fit onto the gpu simultaneously
     """
+
     maintain_input_order: bool  # True if we're returning rows in the order we received them from our input
     num_dependencies: np.ndarray  # number of dependencies for our output slots; indexed by slot idx
     outputs: np.ndarray  # bool per slot; True if this slot is part of our output
@@ -68,11 +70,15 @@ class ExprEvalNode(ExecNode):
     num_output_rows: int
 
     BATCH_SIZE = 64
-    MAX_BUFFERED_ROWS = 512  # maximum number of rows that have been dispatched but not yet returned
+    MAX_BUFFERED_ROWS = 2048  # maximum number of rows that have been dispatched but not yet returned
 
     def __init__(
-        self, row_builder: exprs.RowBuilder, output_exprs: Iterable[exprs.Expr], input_exprs: Iterable[exprs.Expr],
-        input: ExecNode, maintain_input_order: bool = True
+        self,
+        row_builder: exprs.RowBuilder,
+        output_exprs: Iterable[exprs.Expr],
+        input_exprs: Iterable[exprs.Expr],
+        input: ExecNode,
+        maintain_input_order: bool = True,
     ):
         super().__init__(row_builder, output_exprs, input_exprs, input)
         self.maintain_input_order = maintain_input_order
@@ -148,7 +154,9 @@ class ExprEvalNode(ExecNode):
                     self.row_pos_map[id(row)] = self.num_input_rows + idx
             self.num_input_rows += len(batch)
             self.avail_input_rows += len(batch)
-            _logger.debug(f'adding input: batch_size={len(batch)} #input_rows={self.num_input_rows} #avail={self.avail_input_rows}')
+            _logger.debug(
+                f'adding input: batch_size={len(batch)} #input_rows={self.num_input_rows} #avail={self.avail_input_rows}'
+            )
         except StopAsyncIteration:
             self.input_complete = True
             _logger.debug(f'finished input: #input_rows={self.num_input_rows}, #avail={self.avail_input_rows}')
@@ -175,11 +183,11 @@ class ExprEvalNode(ExecNode):
         rows: list[exprs.DataRow]
         if avail_current_batch_rows > num_rows:
             # we only need rows from current_input_batch
-            rows = self.current_input_batch.rows[self.input_row_idx:self.input_row_idx + num_rows]
+            rows = self.current_input_batch.rows[self.input_row_idx : self.input_row_idx + num_rows]
             self.input_row_idx += num_rows
         else:
             # we need rows from both current_/next_input_batch
-            rows = self.current_input_batch.rows[self.input_row_idx:]
+            rows = self.current_input_batch.rows[self.input_row_idx :]
             self.current_input_batch = self.next_input_batch
             self.next_input_batch = None
             self.input_row_idx = 0
@@ -236,6 +244,7 @@ class ExprEvalNode(ExecNode):
         exc_event_aw = asyncio.create_task(self.exc_event.wait(), name='exc_event.wait()')
         input_batch_aw: Optional[asyncio.Task] = None
         completed_aw: Optional[asyncio.Task] = None
+        closed_evaluators = False  # True after calling Evaluator.close()
 
         try:
             while True:
@@ -275,11 +284,12 @@ class ExprEvalNode(ExecNode):
                     assert self.output_buffer.num_rows == 0
                     return
 
-                if self.input_complete and self.avail_input_rows == 0:
+                if self.input_complete and self.avail_input_rows == 0 and not closed_evaluators:
                     # no more input rows to dispatch, but we're still waiting for rows to finish:
                     # close  all slot evaluators to flush queued rows
                     for evaluator in self.slot_evaluators.values():
                         evaluator.close()
+                    closed_evaluators = True
 
                 # we don't have a full batch of rows at this point and need to wait
                 aws = {exc_event_aw}  # always wait for an exception
@@ -335,8 +345,7 @@ class ExprEvalNode(ExecNode):
             first_row = rows[0]
             input_vals = [first_row[idx] for idx in dependency_idxs]
             e = self.row_builder.unique_exprs[slot_with_exc]
-            self.error = excs.ExprEvalError(
-                e, f'expression {e}', first_row.get_exc(e.slot_idx), exc_tb, input_vals, 0)
+            self.error = excs.ExprEvalError(e, f'expression {e}', first_row.get_exc(e.slot_idx), exc_tb, input_vals, 0)
             self.exc_event.set()
             return
 
@@ -391,7 +400,11 @@ class ExprEvalNode(ExecNode):
             _logger.debug(f'Scheduling {len(ready_rows)} rows for slot {slot_idx}')
             self.slot_evaluators[slot_idx].schedule(ready_rows, slot_idx)
 
-    def done_cb(self, t: asyncio.Task) -> None:
+    def register_task(self, t: asyncio.Task) -> None:
+        self.tasks.add(t)
+        t.add_done_callback(self._done_cb)
+
+    def _done_cb(self, t: asyncio.Task) -> None:
         self.tasks.discard(t)
         # end the main loop if we had an unhandled exception
         try:

pixeltable/exec/expr_eval/globals.py

@@ -1,16 +1,18 @@
+from __future__ import annotations
+
 import abc
 import asyncio
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any, Protocol, Optional
+from typing import Any, Optional, Protocol
 
-from pixeltable import exprs
-from pixeltable import func
+from pixeltable import exprs, func
 
 
 @dataclass
 class FnCallArgs:
     """Container for everything needed to execute a FunctionCall against one or more DataRows"""
+
     fn_call: exprs.FunctionCall
     rows: list[exprs.DataRow]
     # single call
@@ -37,16 +39,36 @@ class FnCallArgs:
 
 class Scheduler(abc.ABC):
     """
-    Base class for schedulers. A scheduler executes FunctionCalls against a limited resource pool.
+    Base class for queueing schedulers. A scheduler executes FunctionCalls against a limited resource pool.
 
     Expected behavior:
     - all created tasks must be recorded in dispatcher.tasks
     - schedulers are responsible for aborting execution when a) the task is cancelled or b) when an exception occurred
       elsewhere (indicated by dispatcher.exc_event)
     """
-    @abc.abstractmethod
+
+    @dataclass(frozen=True)
+    class QueueItem:
+        """Container of work items for queueing schedulers"""
+
+        request: FnCallArgs
+        num_retries: int
+
+        def __lt__(self, other: Scheduler.QueueItem) -> bool:
+            # prioritize by number of retries (more retries = higher priority)
+            return self.num_retries > other.num_retries
+
+    resource_pool: str
+    queue: asyncio.PriorityQueue[QueueItem]  # prioritizes retries
+    dispatcher: Dispatcher
+
+    def __init__(self, resource_pool: str, dispatcher: Dispatcher):
+        self.resource_pool = resource_pool
+        self.queue = asyncio.PriorityQueue()
+        self.dispatcher = dispatcher
+
     def submit(self, item: FnCallArgs) -> None:
-        pass
+        self.queue.put_nowait(self.QueueItem(item, 0))
 
     @classmethod
     @abc.abstractmethod
@@ -59,11 +81,11 @@ class Dispatcher(Protocol):
     """
     Row dispatcher used by Evaluators/Schedulers for post-processing after slot materialization and for task management.
 
-    Task management: all tasks need to be recorded in tasks and have done_cb registered with add_done_callback()
+    Task management: all tasks need to be registered via register_task()
     Exceptions: evaluators/schedulers need to check exc_event prior to starting long-running (non-interruptible)
         computations
     """
-    tasks: set[asyncio.Task]
+
     row_builder: exprs.RowBuilder
     exc_event: asyncio.Event
     schedulers: dict[str, Scheduler]  # key: resource pool id
@@ -76,8 +98,8 @@ class Dispatcher(Protocol):
         """Propagates exception in slot_with_exc to all dependent slots and dispatches the rest; does not block"""
         ...
 
-    def done_cb(self, f: asyncio.Task) -> None:
-        """Callback for task completion; does not block"""
+    def register_task(self, f: asyncio.Task) -> None:
+        """Register task with dispatcher for subsequent cleanup; does not block"""
         ...
 
 
@@ -91,6 +113,7 @@ class Evaluator(abc.ABC):
     - evaluators are responsible for aborting execution when a) the task is cancelled or b) when an exception occurred
       elsewhere (indicated by dispatcher.exc_event)
     """
+
     dispatcher: Dispatcher
     is_closed: bool
 
@@ -111,4 +134,3 @@ class Evaluator(abc.ABC):
         """Indicates that there may not be any more rows getting scheduled"""
         self.is_closed = True
         self._close()
-

pixeltable/exec/expr_eval/row_buffer.py

@@ -62,15 +62,14 @@ class RowBuffer:
             return []
         rows: list[exprs.DataRow]
         if self.head_idx + n <= self.size:
-            rows = self.buffer[self.head_idx:self.head_idx + n].tolist()
-            self.buffer[self.head_idx:self.head_idx + n] = None
+            rows = self.buffer[self.head_idx : self.head_idx + n].tolist()
+            self.buffer[self.head_idx : self.head_idx + n] = None
         else:
-            rows = np.concatenate([self.buffer[self.head_idx:], self.buffer[:self.head_idx + n - self.size]]).tolist()
-            self.buffer[self.head_idx:] = None
-            self.buffer[:self.head_idx + n - self.size] = None
+            rows = np.concatenate([self.buffer[self.head_idx :], self.buffer[: self.head_idx + n - self.size]]).tolist()
+            self.buffer[self.head_idx :] = None
+            self.buffer[: self.head_idx + n - self.size] = None
         self.head_pos += n
         self.head_idx = (self.head_idx + n) % self.size
         self.num_rows -= n
         self.num_ready -= n
         return rows
-

pixeltable/exec/expr_eval/schedulers.py

@@ -5,12 +5,12 @@ import datetime
 import inspect
 import logging
 import sys
-from dataclasses import dataclass
-from typing import Optional, Awaitable, Collection
+import time
+from typing import Awaitable, Collection, Optional
 
-from pixeltable import env
-from pixeltable import func
-from .globals import Scheduler, FnCallArgs, Dispatcher
+from pixeltable import env, func
+
+from .globals import Dispatcher, FnCallArgs, Scheduler
 
 _logger = logging.getLogger('pixeltable')
 
@@ -29,19 +29,7 @@ class RateLimitsScheduler(Scheduler):
     TODO:
     - limit the number of in-flight requests based on the open file limit
     """
-    @dataclass(frozen=True)
-    class QueueItem:
-        request: FnCallArgs
-        num_retries: int
-
-        def __lt__(self, other: RateLimitsScheduler.QueueItem) -> bool:
-            # prioritize by number of retries
-            return self.num_retries > other.num_retries
-
-    resource_pool: str
-    queue: asyncio.PriorityQueue[QueueItem]  # prioritizes retries
-    loop_task: asyncio.Task
-    dispatcher: Dispatcher
+
     get_request_resources_param_names: list[str]  # names of parameters of RateLimitsInfo.get_request_resources()
 
     # scheduling-related state
@@ -58,12 +46,9 @@ class RateLimitsScheduler(Scheduler):
     MAX_RETRIES = 10
 
     def __init__(self, resource_pool: str, dispatcher: Dispatcher):
-        self.resource_pool = resource_pool
-        self.queue = asyncio.PriorityQueue()
-        self.dispatcher = dispatcher
-        self.loop_task = asyncio.create_task(self._main_loop())
-        self.dispatcher.tasks.add(self.loop_task)
-        self.loop_task.add_done_callback(self.dispatcher.done_cb)
+        super().__init__(resource_pool, dispatcher)
+        loop_task = asyncio.create_task(self._main_loop())
+        self.dispatcher.register_task(loop_task)
         self.pool_info = None  # initialized in _main_loop by the first request
         self.est_usage = {}
         self.num_in_flight = 0
@@ -79,6 +64,19 @@ class RateLimitsScheduler(Scheduler):
     def submit(self, item: FnCallArgs) -> None:
         self.queue.put_nowait(self.QueueItem(item, 0))
 
+    def _set_pool_info(self) -> None:
+        """Initialize pool_info with the RateLimitsInfo for the resource pool, if available"""
+        if self.pool_info is not None:
+            return
+        self.pool_info = env.Env.get().get_resource_pool_info(self.resource_pool, None)
+        if self.pool_info is None:
+            return
+        assert isinstance(self.pool_info, env.RateLimitsInfo)
+        assert hasattr(self.pool_info, 'get_request_resources')
+        sig = inspect.signature(self.pool_info.get_request_resources)
+        self.get_request_resources_param_names = [p.name for p in sig.parameters.values()]
+        self.est_usage = {r: 0 for r in self._resources}
+
     async def _main_loop(self) -> None:
         item: Optional[RateLimitsScheduler.QueueItem] = None
         while True:
@@ -92,21 +90,15 @@ class RateLimitsScheduler(Scheduler):
                 # wait for a single request to get rate limits
                 _logger.debug(f'initializing rate limits for {self.resource_pool}')
                 await self._exec(item.request, item.num_retries, is_task=False)
+                _logger.debug(f'initialized rate limits for {self.resource_pool}')
                 item = None
                 # if this was the first request, it created the pool_info
                 if self.pool_info is None:
-                    self.pool_info = env.Env.get().get_resource_pool_info(self.resource_pool, None)
-                    if self.pool_info is None:
-                        # we still don't have rate limits, wait for the next request
-                        continue
-                    assert isinstance(self.pool_info, env.RateLimitsInfo)
-                    assert hasattr(self.pool_info, 'get_request_resources')
-                    sig = inspect.signature(self.pool_info.get_request_resources)
-                    self.get_request_resources_param_names = [p.name for p in sig.parameters.values()]
-                    self.est_usage = {r: 0 for r in self._resources}
+                    self._set_pool_info()
                 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]] = []
@@ -149,8 +141,7 @@ class RateLimitsScheduler(Scheduler):
             _logger.debug(f'creating task for {self.resource_pool}')
             self.num_in_flight += 1
             task = asyncio.create_task(self._exec(item.request, item.num_retries, is_task=True))
-            self.dispatcher.tasks.add(task)
-            task.add_done_callback(self.dispatcher.done_cb)
+            self.dispatcher.register_task(task)
             item = None
 
     @property
@@ -166,7 +157,6 @@ 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)
@@ -188,7 +178,9 @@ class RateLimitsScheduler(Scheduler):
             start_ts = datetime.datetime.now(tz=datetime.timezone.utc)
             pxt_fn = request.fn_call.fn
             assert isinstance(pxt_fn, func.CallableFunction)
-            _logger.debug(f'scheduler {self.resource_pool}: start evaluating slot {request.fn_call.slot_idx}, batch_size={len(request.rows)}')
+            _logger.debug(
+                f'scheduler {self.resource_pool}: start evaluating slot {request.fn_call.slot_idx}, batch_size={len(request.rows)}'
+            )
             self.total_requests += 1
             if request.is_batched:
                 batch_result = await pxt_fn.aexec_batch(*request.batch_args, **request.batch_kwargs)
@@ -199,17 +191,25 @@ class RateLimitsScheduler(Scheduler):
                 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(f'scheduler {self.resource_pool}: evaluated slot {request.fn_call.slot_idx} in {end_ts - start_ts}, batch_size={len(request.rows)}')
+            _logger.debug(
+                f'scheduler {self.resource_pool}: evaluated slot {request.fn_call.slot_idx} in {end_ts - start_ts}, batch_size={len(request.rows)}'
+            )
 
             # purge accumulated usage estimate, now that we have a new report
             self.est_usage = {r: 0 for r in self._resources}
 
             self.dispatcher.dispatch(request.rows)
         except Exception as exc:
-            if num_retries < self.MAX_RETRIES and self.pool_info is not None:
+            _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))
                     return
@@ -221,12 +221,140 @@ class RateLimitsScheduler(Scheduler):
                 row.set_exc(request.fn_call.slot_idx, exc)
             self.dispatcher.dispatch_exc(request.rows, request.fn_call.slot_idx, exc_tb)
         finally:
-            _logger.debug(
-                f'Scheduler stats: #requests={self.total_requests}, #retried={self.total_retried}')
+            _logger.debug(f'Scheduler stats: #requests={self.total_requests}, #retried={self.total_retried}')
             if is_task:
                 self.num_in_flight -= 1
                 self.request_completed.set()
 
 
+class RequestRateScheduler(Scheduler):
+    """
+    Scheduler for FunctionCalls with a fixed request rate limit and no runtime resource usage reports.
+
+    Rate limits are supplied in the config, in one of two ways:
+    - resource_pool='request-rate:<endpoint>':
+      * a single rate limit for all calls against that endpoint
+      * in the config: section '<endpoint>', key 'rate_limit'
+    - resource_pool='request-rate:<endpoint>:<model>':
+        * a single rate limit for all calls against that model
+        * in the config: section '<endpoint>.rate_limits', key '<model>'
+    - if no rate limit is found in the config, uses a default of 600 RPM
+
+    TODO:
+    - adaptive rate limiting based on 429 errors
+    """
+
+    secs_per_request: float  # inverted rate limit
+    num_in_flight: int
+    total_requests: int
+    total_retried: int
+
+    TIME_FORMAT = '%H:%M.%S %f'
+    MAX_RETRIES = 10
+    DEFAULT_RATE_LIMIT = 600  # requests per minute
+
+    def __init__(self, resource_pool: str, dispatcher: Dispatcher):
+        super().__init__(resource_pool, dispatcher)
+        loop_task = asyncio.create_task(self._main_loop())
+        self.dispatcher.register_task(loop_task)
+        self.num_in_flight = 0
+        self.total_requests = 0
+        self.total_retried = 0
+
+        # try to get the rate limit from the config
+        elems = resource_pool.split(':')
+        section: str
+        key: str
+        if len(elems) == 2:
+            # resource_pool: request-rate:endpoint
+            _, endpoint = elems
+            section = endpoint
+            key = 'rate_limit'
+        else:
+            # resource_pool: request-rate:endpoint:model
+            assert len(elems) == 3
+            _, endpoint, model = elems
+            section = f'{endpoint}.rate_limits'
+            key = model
+        requests_per_min = env.Env.get().config.get_int_value(key, section=section)
+        requests_per_min = requests_per_min or self.DEFAULT_RATE_LIMIT
+        self.secs_per_request = 1 / (requests_per_min / 60)
+
+    @classmethod
+    def matches(cls, resource_pool: str) -> bool:
+        return resource_pool.startswith('request-rate:')
+
+    async def _main_loop(self) -> None:
+        last_request_ts = 0.0
+        while True:
+            item = await self.queue.get()
+            if item.num_retries > 0:
+                self.total_retried += 1
+            now = time.monotonic()
+            if now - last_request_ts < self.secs_per_request:
+                wait_duration = self.secs_per_request - (now - last_request_ts)
+                _logger.debug(f'waiting for {wait_duration} for {self.resource_pool}')
+                await asyncio.sleep(wait_duration)
+
+            last_request_ts = time.monotonic()
+            if item.num_retries > 0:
+                # the last request encountered some problem: retry it synchronously, to wait for the problem to pass
+                _logger.debug(f'retrying request for {self.resource_pool}: #retries={item.num_retries}')
+                await self._exec(item.request, item.num_retries, is_task=False)
+                _logger.debug(f'retried request for {self.resource_pool}: #retries={item.num_retries}')
+            else:
+                _logger.debug(f'creating task for {self.resource_pool}')
+                self.num_in_flight += 1
+                task = asyncio.create_task(self._exec(item.request, item.num_retries, is_task=True))
+                self.dispatcher.register_task(task)
+
+    async def _exec(self, request: FnCallArgs, 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)
+
+        try:
+            start_ts = datetime.datetime.now(tz=datetime.timezone.utc)
+            pxt_fn = request.fn_call.fn
+            assert isinstance(pxt_fn, func.CallableFunction)
+            _logger.debug(
+                f'scheduler {self.resource_pool}: start evaluating slot {request.fn_call.slot_idx}, batch_size={len(request.rows)}'
+            )
+            self.total_requests += 1
+            if request.is_batched:
+                batch_result = await pxt_fn.aexec_batch(*request.batch_args, **request.batch_kwargs)
+                assert len(batch_result) == len(request.rows)
+                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.row[request.fn_call.slot_idx] = result
+            end_ts = datetime.datetime.now(tz=datetime.timezone.utc)
+            _logger.debug(
+                f'scheduler {self.resource_pool}: evaluated slot {request.fn_call.slot_idx} in {end_ts - start_ts}, batch_size={len(request.rows)}'
+            )
+            self.dispatcher.dispatch(request.rows)
+
+        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))
+                return
+
+            # record the exception
+            _, _, exc_tb = sys.exc_info()
+            for row in request.rows:
+                row.set_exc(request.fn_call.slot_idx, exc)
+            self.dispatcher.dispatch_exc(request.rows, request.fn_call.slot_idx, exc_tb)
+        finally:
+            _logger.debug(
+                f'Scheduler stats: #in-flight={self.num_in_flight} #requests={self.total_requests}, #retried={self.total_retried}'
+            )
+            if is_task:
+                self.num_in_flight -= 1
+
+
 # all concrete Scheduler subclasses that implement matches()
-SCHEDULERS = [RateLimitsScheduler]
+SCHEDULERS = [RateLimitsScheduler, RequestRateScheduler]

pixeltable/exec/in_memory_data_node.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Iterator, Optional, AsyncIterator
+from typing import Any, AsyncIterator, Iterator, Optional
 
 import pixeltable.catalog as catalog
 import pixeltable.exprs as exprs
@@ -10,6 +10,7 @@ from .exec_node import ExecNode
 
 _logger = logging.getLogger('pixeltable')
 
+
 class InMemoryDataNode(ExecNode):
     """
     Outputs in-memory data as a DataRowBatch of a particular table.
@@ -18,6 +19,7 @@ class InMemoryDataNode(ExecNode):
     - with the values provided in the input rows
     - if an input row doesn't provide a value, sets the slot to the column default
     """
+
     tbl: catalog.TableVersion
     input_rows: list[dict[str, Any]]
     start_row_id: int
@@ -27,8 +29,7 @@ class InMemoryDataNode(ExecNode):
     output_exprs: list[exprs.ColumnRef]
 
     def __init__(
-        self, tbl: catalog.TableVersion, rows: list[dict[str, Any]],
-        row_builder: exprs.RowBuilder, start_row_id: int,
+        self, tbl: catalog.TableVersion, rows: list[dict[str, Any]], row_builder: exprs.RowBuilder, start_row_id: int
     ):
         # we materialize the input slots
         output_exprs = list(row_builder.input_exprs)
@@ -43,11 +44,11 @@ class InMemoryDataNode(ExecNode):
         """Create row batch and populate with self.input_rows"""
         user_cols_by_name = {
             col_ref.col.name: exprs.ColumnSlotIdx(col_ref.col, col_ref.slot_idx)
-            for col_ref in self.output_exprs if col_ref.col.name is not None
+            for col_ref in self.output_exprs
+            if col_ref.col.name is not None
         }
         output_cols_by_idx = {
-            col_ref.slot_idx: exprs.ColumnSlotIdx(col_ref.col, col_ref.slot_idx)
-            for col_ref in self.output_exprs
+            col_ref.slot_idx: exprs.ColumnSlotIdx(col_ref.col, col_ref.slot_idx) for col_ref in self.output_exprs
         }
         output_slot_idxs = {e.slot_idx for e in self.output_exprs}
 
@@ -68,7 +69,7 @@ class InMemoryDataNode(ExecNode):
                 input_slot_idxs.add(col_info.slot_idx)
 
             # set the remaining output slots to their default values (presently None)
-            missing_slot_idxs =  output_slot_idxs - input_slot_idxs
+            missing_slot_idxs = output_slot_idxs - input_slot_idxs
             for slot_idx in missing_slot_idxs:
                 col_info = output_cols_by_idx.get(slot_idx)
                 assert col_info is not None

pixeltable/exec/row_update_node.py

@@ -4,11 +4,13 @@ from typing import Any, AsyncIterator
 import pixeltable.catalog as catalog
 import pixeltable.exprs as exprs
 from pixeltable.utils.media_store import MediaStore
+
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
 
 _logger = logging.getLogger('pixeltable')
 
+
 class RowUpdateNode(ExecNode):
     """
     Update individual rows in the input batches, identified by key columns.
@@ -17,9 +19,15 @@ class RowUpdateNode(ExecNode):
     The node assumes that all update dicts contain the same keys, and it populates the slots of the columns present in
     the update list.
     """
+
     def __init__(
-            self, tbl: catalog.TableVersionPath, key_vals_batch: list[tuple], is_rowid_key: bool,
-            col_vals_batch: list[dict[catalog.Column, Any]], row_builder: exprs.RowBuilder, input: ExecNode,
+        self,
+        tbl: catalog.TableVersionPath,
+        key_vals_batch: list[tuple],
+        is_rowid_key: bool,
+        col_vals_batch: list[dict[catalog.Column, Any]],
+        row_builder: exprs.RowBuilder,
+        input: ExecNode,
     ):
         super().__init__(row_builder, [], [], input)
         self.updates = {key_vals: col_vals for key_vals, col_vals in zip(key_vals_batch, col_vals_batch)}
@@ -28,7 +36,8 @@ class RowUpdateNode(ExecNode):
         # retrieve ColumnRefs from the RowBuilder (has slot_idx set)
         all_col_slot_idxs = {
             col_ref.col: col_ref.slot_idx
-            for col_ref in row_builder.unique_exprs if isinstance(col_ref, exprs.ColumnRef)
+            for col_ref in row_builder.unique_exprs
+            if isinstance(col_ref, exprs.ColumnRef)
         }
         self.col_slot_idxs = {col: all_col_slot_idxs[col] for col in col_vals_batch[0].keys()}
         self.key_slot_idxs = {col: all_col_slot_idxs[col] for col in tbl.tbl_version.primary_key_columns()}
@@ -37,8 +46,9 @@ class RowUpdateNode(ExecNode):
     async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         async for batch in self.input:
             for row in batch:
-                key_vals = row.rowid if self.is_rowid_key else \
-                    tuple(row[slot_idx] for slot_idx in self.key_slot_idxs.values())
+                key_vals = (
+                    row.rowid if self.is_rowid_key else tuple(row[slot_idx] for slot_idx in self.key_slot_idxs.values())
+                )
                 if key_vals not in self.updates:
                     continue
                 self.matched_key_vals.add(key_vals)