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

pixeltable/catalog/catalog.py

@@ -189,12 +189,10 @@ class Catalog:
     @classmethod
     def clear(cls) -> None:
         """Remove the instance. Used for testing."""
-        # invalidate all existing instances to force reloading of metadata
-        for tbl_version in cls._instance._tbl_versions.values():
-            # _logger.debug(
-            #     f'Invalidating table version {tbl_version.id}:{tbl_version.effective_version} ({id(tbl_version):x})'
-            # )
-            tbl_version.is_validated = False
+        if cls._instance is not None:
+            # invalidate all existing instances to force reloading of metadata
+            for tbl_version in cls._instance._tbl_versions.values():
+                tbl_version.is_validated = False
         cls._instance = None
 
     def __init__(self) -> None:

pixeltable/catalog/table.py

@@ -183,16 +183,14 @@ class Table(SchemaObject):
 
         return op()
 
-    def _get_views(self, *, recursive: bool = True, include_snapshots: bool = True) -> list['Table']:
+    def _get_views(self, *, recursive: bool = True, mutable_only: bool = False) -> list['Table']:
         cat = catalog.Catalog.get()
         view_ids = cat.get_view_ids(self._id)
         views = [cat.get_table_by_id(id) for id in view_ids]
-        if not include_snapshots:
-            views = [t for t in views if not t._tbl_version_path.is_snapshot()]
+        if mutable_only:
+            views = [t for t in views if t._tbl_version_path.is_mutable()]
         if recursive:
-            views.extend(
-                t for view in views for t in view._get_views(recursive=True, include_snapshots=include_snapshots)
-            )
+            views.extend(t for view in views for t in view._get_views(recursive=True, mutable_only=mutable_only))
         return views
 
     def _df(self) -> 'pxt.dataframe.DataFrame':
@@ -836,21 +834,25 @@ class Table(SchemaObject):
             if_not_exists_ = IfNotExistsParam.validated(if_not_exists, 'if_not_exists')
 
             if isinstance(column, str):
-                col = self._tbl_version_path.get_column(column, include_bases=False)
+                col = self._tbl_version_path.get_column(column)
                 if col is None:
                     if if_not_exists_ == IfNotExistsParam.ERROR:
                         raise excs.Error(f'Column {column!r} unknown')
                     assert if_not_exists_ == IfNotExistsParam.IGNORE
                     return
+                if col.tbl.id != self._tbl_version_path.tbl_id:
+                    raise excs.Error(f'Cannot drop base table column {col.name!r}')
                 col = self._tbl_version.get().cols_by_name[column]
             else:
-                exists = self._tbl_version_path.has_column(column.col, include_bases=False)
+                exists = self._tbl_version_path.has_column(column.col)
                 if not exists:
                     if if_not_exists_ == IfNotExistsParam.ERROR:
                         raise excs.Error(f'Unknown column: {column.col.qualified_name}')
                     assert if_not_exists_ == IfNotExistsParam.IGNORE
                     return
                 col = column.col
+                if col.tbl.id != self._tbl_version_path.tbl_id:
+                    raise excs.Error(f'Cannot drop base table column {col.name!r}')
 
             dependent_user_cols = [c for c in cat.get_column_dependents(col.tbl.id, col.id) if c.name is not None]
             if len(dependent_user_cols) > 0:
@@ -859,13 +861,32 @@ class Table(SchemaObject):
                     f'{", ".join(c.name for c in dependent_user_cols)}'
                 )
 
-            _ = self._get_views(recursive=True, include_snapshots=False)
+            views = self._get_views(recursive=True, mutable_only=True)
+
+            # See if any view predicates depend on this column
+            dependent_views = []
+            for view in views:
+                if view._tbl_version is not None:
+                    predicate = view._tbl_version.get().predicate
+                    if predicate is not None:
+                        for predicate_col in exprs.Expr.get_refd_column_ids(predicate.as_dict()):
+                            if predicate_col.tbl_id == col.tbl.id and predicate_col.col_id == col.id:
+                                dependent_views.append((view, predicate))
+
+            if len(dependent_views) > 0:
+                dependent_views_str = '\n'.join(
+                    f'view: {view._path()}, predicate: {predicate!s}' for view, predicate in dependent_views
+                )
+                raise excs.Error(
+                    f'Cannot drop column `{col.name}` because the following views depend on it:\n{dependent_views_str}'
+                )
+
             # See if this column has a dependent store. We need to look through all stores in all
             # (transitive) views of this table.
             col_handle = col.handle
             dependent_stores = [
                 (view, store)
-                for view in (self, *self._get_views(recursive=True, include_snapshots=False))
+                for view in (self, *views)
                 for store in view._tbl_version.get().external_stores.values()
                 if col_handle in store.get_local_columns()
             ]
@@ -878,6 +899,12 @@ class Table(SchemaObject):
                     f'Cannot drop column `{col.name}` because the following external stores depend on it:\n'
                     f'{", ".join(dependent_store_names)}'
                 )
+            all_columns = self.columns()
+            if len(all_columns) == 1 and col.name == all_columns[0]:
+                raise excs.Error(
+                    f'Cannot drop column `{col.name}` because it is the last remaining column in this table.'
+                    f' Tables must have at least one column.'
+                )
 
             self._tbl_version.get().drop_column(col)
 
@@ -1108,11 +1135,11 @@ class Table(SchemaObject):
         """Resolve a column parameter to a Column object"""
         col: Column = None
         if isinstance(column, str):
-            col = self._tbl_version_path.get_column(column, include_bases=True)
+            col = self._tbl_version_path.get_column(column)
             if col is None:
                 raise excs.Error(f'Column {column!r} unknown')
         elif isinstance(column, ColumnRef):
-            exists = self._tbl_version_path.has_column(column.col, include_bases=True)
+            exists = self._tbl_version_path.has_column(column.col)
             if not exists:
                 raise excs.Error(f'Unknown column: {column.col.qualified_name}')
             col = column.col
@@ -1483,14 +1510,14 @@ class Table(SchemaObject):
                 col_name: str
                 col: Column
                 if isinstance(column, str):
-                    col = self._tbl_version_path.get_column(column, include_bases=True)
+                    col = self._tbl_version_path.get_column(column)
                     if col is None:
                         raise excs.Error(f'Unknown column: {column!r}')
                     col_name = column
                 else:
                     assert isinstance(column, ColumnRef)
                     col = column.col
-                    if not self._tbl_version_path.has_column(col, include_bases=True):
+                    if not self._tbl_version_path.has_column(col):
                         raise excs.Error(f'Unknown column: {col.name!r}')
                     col_name = col.name
                 if not col.is_computed:

pixeltable/catalog/table_version.py

@@ -327,7 +327,7 @@ class TableVersion:
             from .table_version_path import TableVersionPath
 
             # clear out any remaining media files from an aborted previous attempt
-            MediaStore.delete(self.id)
+            MediaStore.get().delete(self.id)
             view_path = TableVersionPath.from_dict(op.load_view_op.view_path)
             plan, _ = Planner.create_view_load_plan(view_path)
             _, row_counts = self.store_tbl.insert_rows(plan, v_min=self.version)
@@ -374,7 +374,7 @@ class TableVersion:
         #     if self.base.get().is_mutable:
         #         self.base.get().mutable_views.remove(TableVersionHandle.create(self))
 
-        MediaStore.delete(self.id)
+        MediaStore.get().delete(self.id)
         FileCache.get().clear(tbl_id=self.id)
         self.store_tbl.drop()
 
@@ -827,14 +827,17 @@ class TableVersion:
 
     def rename_column(self, old_name: str, new_name: str) -> None:
         """Rename a column."""
-        assert self.is_mutable
-        if old_name not in self.cols_by_name:
+        if not self.is_mutable:
+            raise excs.Error(f'Cannot rename column for immutable table {self.name!r}')
+        col = self.path.get_column(old_name)
+        if col is None:
             raise excs.Error(f'Unknown column: {old_name}')
+        if col.tbl.id != self.id:
+            raise excs.Error(f'Cannot rename base table column {col.name!r}')
         if not is_valid_identifier(new_name):
             raise excs.Error(f"Invalid column name: '{new_name}'")
         if new_name in self.cols_by_name:
             raise excs.Error(f'Column {new_name} already exists')
-        col = self.cols_by_name[old_name]
         del self.cols_by_name[old_name]
         col.name = new_name
         self.cols_by_name[new_name] = col
@@ -1024,10 +1027,11 @@ class TableVersion:
                 for el in val:
                     assert isinstance(el, int)
                 continue
-            col = self.path.get_column(col_name, include_bases=False)
+            col = self.path.get_column(col_name)
             if col is None:
-                # TODO: return more informative error if this is trying to update a base column
                 raise excs.Error(f'Column {col_name} unknown')
+            if col.tbl.id != self.id:
+                raise excs.Error(f'Column {col.name!r} is a base table column and cannot be updated')
             if col.is_computed:
                 raise excs.Error(f'Column {col_name} is computed and cannot be updated')
             if col.is_pk and not allow_pk:
@@ -1235,7 +1239,7 @@ class TableVersion:
             )
 
         # delete newly-added data
-        MediaStore.delete(self.id, tbl_version=self.version)
+        MediaStore.get().delete(self.id, tbl_version=self.version)
         conn.execute(sql.delete(self.store_tbl.sa_tbl).where(self.store_tbl.sa_tbl.c.v_min == self.version))
 
         # revert new deletions

pixeltable/catalog/table_version_path.py

@@ -184,13 +184,13 @@ class TableVersionPath:
         cols = self.columns()
         return {col.id: col for col in cols}
 
-    def get_column(self, name: str, include_bases: Optional[bool] = None) -> Optional[Column]:
+    def get_column(self, name: str) -> Optional[Column]:
         """Return the column with the given name, or None if not found"""
         self.refresh_cached_md()
         col = self._cached_tbl_version.cols_by_name.get(name)
         if col is not None:
             return col
-        elif self.base is not None and (include_bases or self._cached_tbl_version.include_base_columns):
+        elif self.base is not None and self._cached_tbl_version.include_base_columns:
             return self.base.get_column(name)
         else:
             return None
@@ -206,10 +206,11 @@ class TableVersionPath:
         else:
             return None
 
-    def has_column(self, col: Column, include_bases: bool = True) -> bool:
+    def has_column(self, col: Column) -> bool:
         """Return True if this table has the given column."""
-        self.refresh_cached_md()
         assert col.tbl is not None
+        self.refresh_cached_md()
+
         if (
             col.tbl.id == self.tbl_version.id
             and col.tbl.effective_version == self.tbl_version.effective_version
@@ -217,7 +218,7 @@ class TableVersionPath:
         ):
             # the column is visible in this table version
             return True
-        elif self.base is not None and include_bases:
+        elif self.base is not None:
             return self.base.has_column(col)
         else:
             return False

pixeltable/config.py

@@ -111,10 +111,19 @@ class Config:
         return default
 
     def get_value(self, key: str, expected_type: type[T], section: str = 'pixeltable') -> Optional[T]:
-        value = self.lookup_env(section, key)  # Try to get from environment first
+        value: Any = self.lookup_env(section, key)  # Try to get from environment first
         # Next try the config file
-        if value is None and section in self.__config_dict and key in self.__config_dict[section]:
-            value = self.__config_dict[section][key]
+        if value is None:
+            # Resolve nested section dicts
+            lookup_elems = [*section.split('.'), key]
+            value = self.__config_dict
+            for el in lookup_elems:
+                if isinstance(value, dict):
+                    if el not in value:
+                        return None
+                    value = value[el]
+                else:
+                    return None
 
         if value is None:
             return None  # Not specified
@@ -155,19 +164,25 @@ KNOWN_CONFIG_OPTIONS = {
     },
     'anthropic': {'api_key': 'Anthropic API key'},
     'bedrock': {'api_key': 'AWS Bedrock API key'},
-    'deepseek': {'api_key': 'Deepseek API key'},
-    'fireworks': {'api_key': 'Fireworks API key'},
-    'gemini': {'api_key': 'Gemini API key'},
-    'groq': {'api_key': 'Groq API key'},
+    'deepseek': {'api_key': 'Deepseek API key', 'rate_limit': 'Rate limit for Deepseek API requests'},
+    'fireworks': {'api_key': 'Fireworks API key', 'rate_limit': 'Rate limit for Fireworks API requests'},
+    'gemini': {'api_key': 'Gemini API key', 'rate_limits': 'Per-model rate limits for Gemini API requests'},
+    'imagen': {'rate_limits': 'Per-model rate limits for Imagen API requests'},
+    'veo': {'rate_limits': 'Per-model rate limits for Veo API requests'},
+    'groq': {'api_key': 'Groq API key', 'rate_limit': 'Rate limit for Groq API requests'},
     'label_studio': {'api_key': 'Label Studio API key', 'url': 'Label Studio server URL'},
-    'mistral': {'api_key': 'Mistral API key'},
+    'mistral': {'api_key': 'Mistral API key', 'rate_limit': 'Rate limit for Mistral API requests'},
     'openai': {
         'api_key': 'OpenAI API key',
         'base_url': 'OpenAI API base URL',
         'api_version': 'API version if using Azure OpenAI',
+        'rate_limits': 'Per-model rate limits for OpenAI API requests',
     },
     'replicate': {'api_token': 'Replicate API token'},
-    'together': {'api_key': 'Together API key'},
+    'together': {
+        'api_key': 'Together API key',
+        'rate_limits': 'Per-model category rate limits for Together API requests',
+    },
     'pypi': {'api_key': 'PyPI API key (for internal use only)'},
 }
 

pixeltable/dataframe.py

@@ -795,19 +795,19 @@ class DataFrame:
         assert len(col_refs) > 0 and len(joined_tbls) >= 2
         for col_ref in col_refs:
             # identify the referenced column by name in 'other'
-            rhs_col = other.get_column(col_ref.col.name, include_bases=True)
+            rhs_col = other.get_column(col_ref.col.name)
             if rhs_col is None:
                 raise excs.Error(f"'on': column {col_ref.col.name!r} not found in joined table")
             rhs_col_ref = exprs.ColumnRef(rhs_col)
 
             lhs_col_ref: Optional[exprs.ColumnRef] = None
-            if any(tbl.has_column(col_ref.col, include_bases=True) for tbl in self._from_clause.tbls):
+            if any(tbl.has_column(col_ref.col) for tbl in self._from_clause.tbls):
                 # col_ref comes from the existing from_clause, we use that directly
                 lhs_col_ref = col_ref
             else:
                 # col_ref comes from other, we need to look for a match in the existing from_clause by name
                 for tbl in self._from_clause.tbls:
-                    col = tbl.get_column(col_ref.col.name, include_bases=True)
+                    col = tbl.get_column(col_ref.col.name)
                     if col is None:
                         continue
                     if lhs_col_ref is not None:

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
@@ -101,6 +100,8 @@ 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)
@@ -246,7 +247,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 +486,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
@@ -557,6 +558,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 +574,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 +752,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 +812,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