pixeltable 0.4.7__py3-none-any.whl → 0.4.9__py3-none-any.whl

pixeltable/env.py

@@ -15,7 +15,6 @@ import sys
 import threading
 import types
 import typing
-import uuid
 import warnings
 from contextlib import contextmanager
 from dataclasses import dataclass, field
@@ -28,6 +27,7 @@ import nest_asyncio  # type: ignore[import-untyped]
 import pixeltable_pgserver
 import sqlalchemy as sql
 from pillow_heif import register_heif_opener  # type: ignore[import-untyped]
+from tenacity import retry, stop_after_attempt, wait_exponential_jitter
 from tqdm import TqdmWarning
 
 from pixeltable import exceptions as excs
@@ -101,12 +101,18 @@ class Env:
     def _init_env(cls, reinit_db: bool = False) -> None:
         assert not cls.__initializing, 'Circular env initialization detected.'
         cls.__initializing = True
+        if cls._instance is not None:
+            cls._instance._clean_up()
         cls._instance = None
         env = Env()
-        env._set_up(reinit_db=reinit_db)
-        env._upgrade_metadata()
-        cls._instance = env
-        cls.__initializing = False
+        try:
+            env._set_up(reinit_db=reinit_db)
+            env._upgrade_metadata()
+            cls._instance = env
+        finally:
+            # Reset the initializing flag, even if setup fails.
+            # This prevents the environment from being left in a broken state.
+            cls.__initializing = False
 
     def __init__(self) -> None:
         assert self._instance is None, 'Env is a singleton; use Env.get() to access the instance'
@@ -246,7 +252,7 @@ class Env:
         if self._current_conn is None:
             assert self._current_session is None
             try:
-                self._current_isolation_level = 'SERIALIZABLE' if for_write else 'REPEATABLE_READ'
+                self._current_isolation_level = 'SERIALIZABLE'
                 with (
                     self.engine.connect().execution_options(isolation_level=self._current_isolation_level) as conn,
                     sql.orm.Session(conn) as session,
@@ -485,7 +491,7 @@ class Env:
                 raise excs.Error(error)
             self._logger.info(f'Using database at: {self.db_url}')
         else:
-            self._db_name = os.environ.get('PIXELTABLE_DB', 'pixeltable')
+            self._db_name = config.get_string_value('db') or 'pixeltable'
             self._pgdata_dir = Path(os.environ.get('PIXELTABLE_PGDATA', str(Config.get().home / 'pgdata')))
             # cleanup_mode=None will leave the postgres process running after Python exits
             # cleanup_mode='stop' will terminate the postgres process when Python exits
@@ -499,14 +505,24 @@ class Env:
         assert self._db_url is not None
         assert self._db_name is not None
 
+    @retry(
+        stop=stop_after_attempt(3),  # Stop after 3 attempts
+        wait=wait_exponential_jitter(initial=0.2, max=1.0, jitter=0.2),  # Exponential backoff with jitter
+    )
     def _init_metadata(self) -> None:
         """
         Create pixeltable metadata tables and system metadata.
         This is an idempotent operation.
+
+        Retry logic handles race conditions when multiple Pixeltable processes
+        attempt to initialize metadata tables simultaneously. The first process may succeed
+        in creating tables while others encounter database constraints (e.g., "table already exists").
+        Exponential backoff with jitter reduces contention between competing processes.
         """
         assert self._sa_engine is not None
         from pixeltable import metadata
 
+        self._logger.debug('Creating pixeltable metadata')
         metadata.schema.base_metadata.create_all(self._sa_engine, checkfirst=True)
         metadata.create_system_info(self._sa_engine)
 
@@ -557,6 +573,14 @@ class Env:
         finally:
             engine.dispose()
 
+    def _pgserver_terminate_connections_stmt(self) -> str:
+        return f"""
+                SELECT pg_terminate_backend(pg_stat_activity.pid)
+                FROM pg_stat_activity
+                WHERE pg_stat_activity.datname = '{self._db_name}'
+                AND pid <> pg_backend_pid()
+            """
+
     def _drop_store_db(self) -> None:
         assert self._db_name is not None
         engine = sql.create_engine(self._dbms.default_system_db_url(), future=True, isolation_level='AUTOCOMMIT')
@@ -565,13 +589,7 @@ class Env:
             with engine.begin() as conn:
                 # terminate active connections
                 if self._db_server is not None:
-                    stmt = f"""
-                        SELECT pg_terminate_backend(pg_stat_activity.pid)
-                        FROM pg_stat_activity
-                        WHERE pg_stat_activity.datname = '{self._db_name}'
-                        AND pid <> pg_backend_pid()
-                    """
-                    conn.execute(sql.text(stmt))
+                    conn.execute(sql.text(self._pgserver_terminate_connections_stmt()))
                 # drop db
                 stmt = self._dbms.drop_db_stmt(preparer.quote(self._db_name))
                 conn.execute(sql.text(stmt))
@@ -749,12 +767,6 @@ class Env:
             else:
                 os.remove(path)
 
-    def num_tmp_files(self) -> int:
-        return len(glob.glob(f'{self._tmp_dir}/*'))
-
-    def create_tmp_path(self, extension: str = '') -> Path:
-        return self._tmp_dir / f'{uuid.uuid4()}{extension}'
-
     # def get_resource_pool_info(self, pool_id: str, pool_info_cls: Optional[Type[T]]) -> T:
     def get_resource_pool_info(self, pool_id: str, make_pool_info: Optional[Callable[[], T]] = None) -> T:
         """Returns the info object for the given id, creating it if necessary."""
@@ -815,6 +827,63 @@ class Env:
         except Exception as exc:
             raise excs.Error(f'Failed to load spaCy model: {spacy_model}') from exc
 
+    def _clean_up(self) -> None:
+        """
+        Internal cleanup method that properly closes all resources and resets state.
+        This is called before destroying the singleton instance.
+        """
+        assert self._current_session is None
+        assert self._current_conn is None
+
+        # Stop HTTP server
+        if self._httpd is not None:
+            try:
+                self._httpd.shutdown()
+                self._httpd.server_close()
+            except Exception as e:
+                _logger.warning(f'Error stopping HTTP server: {e}')
+
+        # First terminate all connections to the database
+        if self._db_server is not None:
+            assert self._dbms is not None
+            assert self._db_name is not None
+            try:
+                temp_engine = sql.create_engine(self._dbms.default_system_db_url(), isolation_level='AUTOCOMMIT')
+                try:
+                    with temp_engine.begin() as conn:
+                        conn.execute(sql.text(self._pgserver_terminate_connections_stmt()))
+                        _logger.info(f"Terminated all connections to database '{self._db_name}'")
+                except Exception as e:
+                    _logger.warning(f'Error terminating database connections: {e}')
+                finally:
+                    temp_engine.dispose()
+            except Exception as e:
+                _logger.warning(f'Error stopping database server: {e}')
+
+        # Dispose of SQLAlchemy engine (after stopping db server)
+        if self._sa_engine is not None:
+            try:
+                self._sa_engine.dispose()
+            except Exception as e:
+                _logger.warning(f'Error disposing engine: {e}')
+
+        # Close event loop
+        if self._event_loop is not None:
+            try:
+                if self._event_loop.is_running():
+                    self._event_loop.stop()
+                self._event_loop.close()
+            except Exception as e:
+                _logger.warning(f'Error closing event loop: {e}')
+
+        # Remove logging handlers
+        for handler in self._logger.handlers[:]:
+            try:
+                handler.close()
+                self._logger.removeHandler(handler)
+            except Exception as e:
+                _logger.warning(f'Error removing handler: {e}')
+
 
 def register_client(name: str) -> Callable:
     """Decorator that registers a third-party API client for use by Pixeltable.

pixeltable/exec/aggregation_node.py

@@ -103,6 +103,6 @@ class AggregationNode(ExecNode):
             self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
             self.output_batch.add_row(prev_row)
 
-        self.output_batch.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
+        self.output_batch.flush_imgs(None, self.row_builder.stored_img_cols, self.flushed_img_slots)
         _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(self.output_batch.rows)} rows')
         yield self.output_batch

pixeltable/exec/cache_prefetch_node.py

@@ -12,8 +12,9 @@ from pathlib import Path
 from typing import Any, AsyncIterator, Iterator, Optional
 from uuid import UUID
 
-from pixeltable import env, exceptions as excs, exprs
+from pixeltable import exceptions as excs, exprs
 from pixeltable.utils.filecache import FileCache
+from pixeltable.utils.media_store import TempStore
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -219,7 +220,7 @@ class CachePrefetchNode(ExecNode):
             self.in_flight_requests[f] = url
 
     def __fetch_url(self, url: str) -> tuple[Optional[Path], Optional[Exception]]:
-        """Fetches a remote URL into Env.tmp_dir and returns its path"""
+        """Fetches a remote URL into the TempStore and returns its path"""
         _logger.debug(f'fetching url={url} thread_name={threading.current_thread().name}')
         parsed = urllib.parse.urlparse(url)
         # Use len(parsed.scheme) > 1 here to ensure we're not being passed
@@ -230,7 +231,7 @@ class CachePrefetchNode(ExecNode):
         if parsed.path:
             p = Path(urllib.parse.unquote(urllib.request.url2pathname(parsed.path)))
             extension = p.suffix
-        tmp_path = env.Env.get().create_tmp_path(extension=extension)
+        tmp_path = TempStore.create_path(extension=extension)
         try:
             _logger.debug(f'Downloading {url} to {tmp_path}')
             if parsed.scheme == 's3':

pixeltable/exec/exec_node.py

@@ -20,7 +20,6 @@ class ExecNode(abc.ABC):
     row_builder: exprs.RowBuilder
     input: Optional[ExecNode]
     flushed_img_slots: list[int]  # idxs of image slots of our output_exprs dependencies
-    stored_img_cols: list[exprs.ColumnSlotIdx]
     ctx: Optional[ExecContext]
 
     def __init__(
@@ -40,7 +39,6 @@ class ExecNode(abc.ABC):
         self.flushed_img_slots = [
             e.slot_idx for e in output_dependencies if e.col_type.is_image_type() and e.slot_idx not in output_slot_idxs
         ]
-        self.stored_img_cols = []
         self.ctx = None  # all nodes of a tree share the same context
 
     def set_ctx(self, ctx: ExecContext) -> None:
@@ -48,12 +46,6 @@ class ExecNode(abc.ABC):
         if self.input is not None:
             self.input.set_ctx(ctx)
 
-    def set_stored_img_cols(self, stored_img_cols: list[exprs.ColumnSlotIdx]) -> None:
-        self.stored_img_cols = stored_img_cols
-        # propagate batch size to the source
-        if self.input is not None:
-            self.input.set_stored_img_cols(stored_img_cols)
-
     @abc.abstractmethod
     def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         pass

pixeltable/exec/expr_eval/globals.py

@@ -56,6 +56,7 @@ class Scheduler(abc.ABC):
         request: FnCallArgs
         num_retries: int
         exec_ctx: ExecCtx
+        retry_after: Optional[float] = None  # time.monotonic()
 
         def __lt__(self, other: Scheduler.QueueItem) -> bool:
             # prioritize by number of retries (more retries = higher priority)

pixeltable/exec/expr_eval/schedulers.py

@@ -270,6 +270,7 @@ class RequestRateScheduler(Scheduler):
     num_in_flight: int
     total_requests: int
     total_retried: int
+    total_errors: int
 
     TIME_FORMAT = '%H:%M.%S %f'
     MAX_RETRIES = 3
@@ -294,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(':')
@@ -312,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
@@ -325,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)
 
@@ -372,15 +379,20 @@ class RequestRateScheduler(Scheduler):
 
         except Exception as exc:
             _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}')
-                await asyncio.sleep(retry_delay)
-                self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx))
+                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)
@@ -388,7 +400,7 @@ 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

pixeltable/exec/in_memory_data_node.py

@@ -2,7 +2,7 @@ import logging
 from typing import Any, AsyncIterator, Optional
 
 from pixeltable import catalog, exprs
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.media_store import TempStore
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -67,8 +67,7 @@ class InMemoryDataNode(ExecNode):
                 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
-                    assert col.tbl.id == self.tbl.id
-                    filepath, _ = MediaStore.save_media_object(val, col, format=None)
+                    filepath, _ = TempStore.save_media_object(val, col, format=None)
                     output_row[col_info.slot_idx] = str(filepath)
                 else:
                     output_row[col_info.slot_idx] = val

pixeltable/exprs/data_row.py

@@ -14,7 +14,7 @@ import PIL.Image
 import sqlalchemy as sql
 
 from pixeltable import catalog, env
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.media_store import MediaStore, TempStore
 
 
 class DataRow:
@@ -270,7 +270,7 @@ class DataRow:
                     # Default to JPEG unless the image has a transparency layer (which isn't supported by JPEG).
                     # In that case, use WebP instead.
                     format = 'webp' if image.has_transparency_data else 'jpeg'
-                filepath, url = MediaStore.save_media_object(image, col, format=format)
+                filepath, url = MediaStore.get().save_media_object(image, col, format=format)
                 self.file_paths[index] = str(filepath)
                 self.file_urls[index] = url
             else:
@@ -282,16 +282,16 @@ class DataRow:
         self.vals[index] = None
 
     def move_tmp_media_file(self, index: int, col: catalog.Column) -> None:
-        """If a media url refers to data in a temporary file, move the data to the MediaStore"""
+        """If a media url refers to data in a temporary file, move the data to a MediaStore"""
         if self.file_urls[index] is None:
             return
         assert self.excs[index] is None
         assert col.col_type.is_media_type()
-        src_path = MediaStore.resolve_tmp_url(self.file_urls[index])
+        src_path = TempStore.resolve_url(self.file_urls[index])
         if src_path is None:
             # The media url does not point to a temporary file, leave it as is
             return
-        new_file_url = MediaStore.relocate_local_media_file(src_path, col)
+        new_file_url = MediaStore.get().relocate_local_media_file(src_path, col)
         self.file_urls[index] = new_file_url
 
     @property

pixeltable/exprs/function_call.py

@@ -115,6 +115,7 @@ class FunctionCall(Expr):
         self._validation_error = validation_error
 
         if validation_error is not None:
+            self.bound_idxs = {}
             self.resource_pool = None
             return
 
@@ -300,8 +301,16 @@ class FunctionCall(Expr):
         """
         res = super().substitute(spec)
         assert res is self
-        self.return_type = self.fn.call_return_type(self.bound_args)
-        self.col_type = self.return_type
+        if self.is_valid:
+            # If this FunctionCall is valid, re-evaluate the call_return_type of the substituted expression. If the
+            # FunctionCall is not valid, it isn't safe to do this. (Really we should be asserting that it *is* valid,
+            # but we still need to be able to do substitutions on invalid FunctionCalls, because loading an
+            # EmbeddingIndex from the db involves reconstructing the requisite (substituted) FunctionCalls. We could
+            # fix this by separately persisting the FunctionCall instances held by EmbeddingIndex to the db. That's
+            # probably a good idea, but it's also probably not urgent, since it only affects Functions that have a
+            # conditional_return_type implemented.)
+            self.return_type = self.fn.call_return_type(self.bound_args)
+            self.col_type = self.return_type
         return self
 
     def update(self, data_row: DataRow) -> None:
@@ -480,25 +489,54 @@ class FunctionCall(Expr):
             ).strip()
         else:
             # Evaluate the call_return_type as defined in the current codebase.
-            call_return_type = resolved_fn.call_return_type(bound_args)
-            if return_type is None:
-                # Schema versions prior to 25 did not store the return_type in metadata, and there is no obvious way to
-                # infer it during DB migration, so we might encounter a stored return_type of None. In that case, we use
-                # the call_return_type that we just inferred (which matches the deserialization behavior prior to
-                # version 25).
-                return_type = call_return_type
-            elif not return_type.is_supertype_of(call_return_type, ignore_nullable=True):
-                # There is a return_type stored in metadata (schema version >= 25),
-                # and the stored return_type of the UDF call doesn't match the column type of the FunctionCall.
-                validation_error = dedent(
-                    f"""
-                    The return type stored in the database for a UDF call to {fn.self_path!r} no longer
-                    matches its return type as currently defined in the code. This probably means that the
-                    code for {fn.self_path!r} has changed in a backward-incompatible way.
-                    Return type of UDF call in the database: {return_type}
-                    Return type of UDF as currently defined in code: {call_return_type}
-                    """
-                ).strip()
+            call_return_type: Optional[ts.ColumnType] = None
+
+            if isinstance(resolved_fn, func.ExprTemplateFunction) and not resolved_fn.template.expr.is_valid:
+                # The FunctionCall is based on an ExprTemplateFunction, but the template expression is not valid
+                # (because it in turn contains an invalid FunctionCall). In this case, inherit the validation error
+                # from the template expression.
+                validation_error = resolved_fn.template.expr.validation_error
+            else:
+                try:
+                    call_return_type = resolved_fn.call_return_type(bound_args)
+                except ImportError as exc:
+                    validation_error = dedent(
+                        f"""
+                        A UDF call to {fn.self_path!r} could not be fully resolved, because a module required
+                        by the UDF could not be imported:
+                        {exc}
+                        """
+                    )
+
+            assert (call_return_type is None) != (validation_error is None)
+
+            if call_return_type is None and return_type is None:
+                # Schema versions prior to 25 did not store the return_type in metadata, and there is no obvious
+                # way to infer it during DB migration, so we might encounter a stored return_type of None. If the
+                # resolution of call_return_type also fails, then we're out of luck; we have no choice but to
+                # fail-fast.
+                raise excs.Error(validation_error)
+
+            if call_return_type is not None:
+                # call_return_type resolution succeeded.
+                if return_type is None:
+                    # Schema versions prior to 25 did not store the return_type in metadata (as mentioned above), so
+                    # fall back on the call_return_type.
+                    return_type = call_return_type
+                elif not return_type.is_supertype_of(call_return_type, ignore_nullable=True):
+                    # There is a return_type stored in metadata (schema version >= 25),
+                    # and the stored return_type of the UDF call doesn't match the column type of the FunctionCall.
+                    validation_error = dedent(
+                        f"""
+                        The return type stored in the database for a UDF call to {fn.self_path!r} no longer
+                        matches its return type as currently defined in the code. This probably means that the
+                        code for {fn.self_path!r} has changed in a backward-incompatible way.
+                        Return type of UDF call in the database: {return_type}
+                        Return type of UDF as currently defined in code: {call_return_type}
+                        """
+                    ).strip()
+
+        assert return_type is not None  # Guaranteed by the above logic.
 
         fn_call = cls(
             resolved_fn,

pixeltable/exprs/row_builder.py

@@ -86,6 +86,8 @@ class RowBuilder:
     img_slot_idxs: list[int]  # Indices of image slots
     media_slot_idxs: list[int]  # Indices of non-image media slots
     array_slot_idxs: list[int]  # Indices of array slots
+    stored_img_cols: list[exprs.ColumnSlotIdx]
+    stored_media_cols: list[exprs.ColumnSlotIdx]
 
     @dataclass
     class EvalCtx:
@@ -112,6 +114,8 @@ class RowBuilder:
         """
         self.unique_exprs: ExprSet[Expr] = ExprSet()  # dependencies precede their dependents
         self.next_slot_idx = 0
+        self.stored_img_cols = []
+        self.stored_media_cols = []
 
         # record input and output exprs; make copies to avoid reusing execution state
         unique_input_exprs = [self._record_unique_expr(e.copy(), recursive=False) for e in input_exprs]
@@ -246,11 +250,13 @@ class RowBuilder:
     def add_table_column(self, col: catalog.Column, slot_idx: int) -> None:
         """Record a column that is part of the table row"""
         assert self.tbl is not None
-        self.table_columns.append(ColumnSlotIdx(col, slot_idx))
-
-    def output_slot_idxs(self) -> list[ColumnSlotIdx]:
-        """Return ColumnSlotIdx for output columns"""
-        return self.table_columns
+        assert col.is_stored
+        info = ColumnSlotIdx(col, slot_idx)
+        self.table_columns.append(info)
+        if col.col_type.is_media_type():
+            self.stored_media_cols.append(info)
+            if col.col_type.is_image_type():
+                self.stored_img_cols.append(info)
 
     @property
     def num_materialized(self) -> int:

pixeltable/func/expr_template_function.py

@@ -85,13 +85,16 @@ class ExprTemplateFunction(Function):
         conditional_return_type).
         """
         assert not self.is_polymorphic
-        template = self.template
         with_defaults = bound_args.copy()
         with_defaults.update(
-            {param_name: default for param_name, default in template.defaults.items() if param_name not in bound_args}
+            {
+                param_name: default
+                for param_name, default in self.template.defaults.items()
+                if param_name not in bound_args
+            }
         )
         substituted_expr = self.template.expr.copy().substitute(
-            {template.param_exprs[name]: expr for name, expr in with_defaults.items()}
+            {self.template.param_exprs[name]: expr for name, expr in with_defaults.items()}
         )
         return substituted_expr.col_type
 

pixeltable/functions/__init__.py

@@ -26,6 +26,8 @@ from . import (
     video,
     vision,
     whisper,
+    whisperx,
+    yolox,
 )
 from .globals import count, map, max, mean, min, sum
 

pixeltable/functions/anthropic.py

@@ -132,8 +132,7 @@ class AnthropicRateLimitsInfo(env.RateLimitsInfo):
         should_retry_str = exc.response.headers.get('x-should-retry', '')
         if should_retry_str.lower() != 'true':
             return None
-        retry_after_str = exc.response.headers.get('retry-after', '1')
-        return int(retry_after_str)
+        return super().get_retry_delay(exc)
 
 
 @pxt.udf

pixeltable/functions/deepseek.py

@@ -26,7 +26,7 @@ def _deepseek_client() -> 'openai.AsyncOpenAI':
     return env.Env.get().get_client('deepseek')
 
 
-@pxt.udf
+@pxt.udf(resource_pool='request-rate:deepseek')
 async def chat_completions(
     messages: list,
     *,
@@ -43,6 +43,10 @@ async def chat_completions(
 
     Deepseek uses the OpenAI SDK, so you will need to install the `openai` package to use this UDF.
 
+    Request throttling:
+    Applies the rate limit set in the config (section `deepseek`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install openai`

pixeltable/functions/gemini.py

@@ -14,6 +14,7 @@ import PIL.Image
 
 import pixeltable as pxt
 from pixeltable import env, exceptions as excs, exprs
+from pixeltable.utils.media_store import TempStore
 
 if TYPE_CHECKING:
     from google import genai
@@ -39,7 +40,7 @@ async def generate_content(
     <https://ai.google.dev/gemini-api/docs/text-generation>
 
     Request throttling:
-    Applies the rate limit set in the config (section `gemini`, key `rate_limit`). If no rate
+    Applies the rate limit set in the config (section `gemini.rate_limits`; use the model id as the key). If no rate
     limit is configured, uses a default of 600 RPM.
 
     __Requirements:__
@@ -126,6 +127,10 @@ async def generate_images(prompt: str, *, model: str, config: Optional[dict] = N
     Generates images based on a text description and configuration. For additional details, see:
     <https://ai.google.dev/gemini-api/docs/image-generation>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `imagen.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install google-genai`
@@ -167,6 +172,10 @@ async def generate_videos(
     Generates videos based on a text description and configuration. For additional details, see:
     <https://ai.google.dev/gemini-api/docs/video-generation>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `veo.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install google-genai`
@@ -215,7 +224,7 @@ async def generate_videos(
     assert video_bytes is not None
 
     # Create a temporary file to store the video bytes
-    output_path = env.Env.get().create_tmp_path('.mp4')
+    output_path = TempStore.create_path(extension='.mp4')
     Path(output_path).write_bytes(video_bytes)
     return str(output_path)