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

pixeltable/env.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from abc import abstractmethod
 import datetime
 import glob
 import http.server
@@ -16,9 +15,11 @@ import sys
 import threading
 import uuid
 import warnings
+from abc import abstractmethod
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Type
+from sys import stdout
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
 from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
 import pixeltable_pgserver
@@ -28,6 +29,7 @@ from tqdm import TqdmWarning
 
 import pixeltable.exceptions as excs
 from pixeltable import metadata
+from pixeltable.utils.console_output import ConsoleLogger, ConsoleMessageFilter, ConsoleOutputHandler, map_level
 from pixeltable.utils.http_server import make_server
 
 if TYPE_CHECKING:
@@ -67,6 +69,7 @@ class Env:
     _httpd: Optional[http.server.HTTPServer]
     _http_address: Optional[str]
     _logger: logging.Logger
+    _console_logger: ConsoleLogger
     _default_log_level: int
     _logfilename: Optional[str]
     _log_to_stdout: bool
@@ -92,6 +95,8 @@ class Env:
         cls._instance = env
 
     def __init__(self):
+        assert self._instance is None, 'Env is a singleton; use Env.get() to access the instance'
+
         self._home = None
         self._media_dir = None  # computed media files
         self._file_cache_dir = None  # cached media files with external URL
@@ -231,6 +236,10 @@ class Env:
         else:
             return False
 
+    @property
+    def console_logger(self) -> ConsoleLogger:
+        return self._console_logger
+
     def _set_up(self, echo: bool = False, reinit_db: bool = False) -> None:
         if self._initialized:
             return
@@ -288,6 +297,14 @@ class Env:
             warnings.simplefilter('ignore', category=UserWarning)
             warnings.simplefilter('ignore', category=FutureWarning)
 
+        # Set verbose level for user visible console messages
+        verbosity = map_level(self._config.get_int_value('verbosity'))
+        stdout_handler = ConsoleOutputHandler(stream=stdout)
+        stdout_handler.setLevel(verbosity)
+        stdout_handler.addFilter(ConsoleMessageFilter())
+        self._logger.addHandler(stdout_handler)
+        self._console_logger = ConsoleLogger(self._logger)
+
         # configure _logger to log to a file
         self._logfilename = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') + '.log'
         fh = logging.FileHandler(self._log_dir / self._logfilename, mode='w')
@@ -358,10 +375,11 @@ class Env:
 
         if create_db:
             from pixeltable.metadata import schema
+
             schema.base_metadata.create_all(self._sa_engine)
             metadata.create_system_info(self._sa_engine)
 
-        print(f'Connected to Pixeltable database at: {self.db_url}')
+        self.console_logger.info(f'Connected to Pixeltable database at: {self.db_url}')
 
         # we now have a home directory and db; start other services
         self._set_up_runtime()
@@ -370,11 +388,7 @@ class Env:
     def _create_engine(self, time_zone_name: Optional[str], echo: bool = False) -> None:
         connect_args = {} if time_zone_name is None else {'options': f'-c timezone={time_zone_name}'}
         self._sa_engine = sql.create_engine(
-            self.db_url,
-            echo=echo,
-            future=True,
-            isolation_level='REPEATABLE READ',
-            connect_args=connect_args,
+            self.db_url, echo=echo, future=True, isolation_level='REPEATABLE READ', connect_args=connect_args
         )
         self._logger.info(f'Created SQLAlchemy engine at: {self.db_url}')
         with self.engine.begin() as conn:
@@ -407,7 +421,7 @@ class Env:
             with engine.begin() as conn:
                 # use C collation to get standard C/Python-style sorting
                 stmt = (
-                    f"CREATE DATABASE {preparer.quote(self._db_name)} "
+                    f'CREATE DATABASE {preparer.quote(self._db_name)} '
                     "ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C' TEMPLATE template0"
                 )
                 conn.execute(sql.text(stmt))
@@ -431,12 +445,12 @@ class Env:
         try:
             with engine.begin() as conn:
                 # terminate active connections
-                stmt = (f"""
+                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))
                 # drop db
                 stmt = f'DROP DATABASE {preparer.quote(self._db_name)}'
@@ -546,7 +560,7 @@ class Env:
             is_installed = False
         self.__optional_packages[package_name] = PackageInfo(
             is_installed=is_installed,
-            library_name=library_name or package_name  # defaults to package_name unless specified otherwise
+            library_name=library_name or package_name,  # defaults to package_name unless specified otherwise
         )
 
     def require_package(self, package_name: str, min_version: Optional[list[int]] = None) -> None:
@@ -592,6 +606,7 @@ class Env:
         """
         import spacy
         from spacy.cli.download import get_model_filename
+
         spacy_model = 'en_core_web_sm'
         spacy_model_version = '3.7.1'
         filename = get_model_filename(spacy_model, spacy_model_version, sdist=False)
@@ -609,7 +624,7 @@ class Env:
             self._logger.warn(f'Failed to load spaCy model: {spacy_model}', exc_info=exc)
             warnings.warn(
                 f"Failed to load spaCy model '{spacy_model}'. spaCy features will not be available.",
-                excs.PixeltableWarning
+                excs.PixeltableWarning,
             )
             self.__optional_packages['spacy'].is_installed = False
 
@@ -619,8 +634,7 @@ class Env:
     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, 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."""
         info = self._resource_pool_info.get(pool_id)
@@ -690,6 +704,7 @@ def register_client(name: str) -> Callable:
     Args:
         - name (str): The name of the API client (e.g., 'openai' or 'label-studio').
     """
+
     def decorator(fn: Callable) -> None:
         global _registered_clients
         sig = inspect.signature(fn)
@@ -704,6 +719,7 @@ class Config:
     The (global) Pixeltable configuration, as loaded from `config.toml`. Provides methods for retrieving
     configuration values, which can be set in the config file or as environment variables.
     """
+
     __config: dict[str, Any]
 
     @classmethod
@@ -733,12 +749,7 @@ class Config:
         free_disk_space_bytes = shutil.disk_usage(config_path.parent).free
         # Default cache size is 1/5 of free disk space
         file_cache_size_g = free_disk_space_bytes / 5 / (1 << 30)
-        return {
-            'pixeltable': {
-                'file_cache_size_g': round(file_cache_size_g, 1),
-                'hide_warnings': False,
-            }
-        }
+        return {'pixeltable': {'file_cache_size_g': round(file_cache_size_g, 1), 'hide_warnings': False}}
 
     def __init__(self, config: dict[str, Any]) -> None:
         self.__config = config
@@ -823,7 +834,9 @@ class RateLimitsInfo:
             self.resource_limits = {k: RateLimitInfo(k, now, *v) for k, v in kwargs.items() if v is not None}
             # TODO: remove
             for info in self.resource_limits.values():
-                _logger.debug(f'Init {info.resource} rate limit: rem={info.remaining} reset={info.reset_at.strftime(TIME_FORMAT)} delta={(info.reset_at - now).total_seconds()}')
+                _logger.debug(
+                    f'Init {info.resource} rate limit: rem={info.remaining} reset={info.reset_at.strftime(TIME_FORMAT)} delta={(info.reset_at - now).total_seconds()}'
+                )
         else:
             for k, v in kwargs.items():
                 if v is not None:
@@ -838,6 +851,7 @@ class RateLimitsInfo:
 @dataclass
 class RateLimitInfo:
     """Container for rate limit-related information for a single resource."""
+
     resource: str
     recorded_at: datetime.datetime
     limit: int
@@ -854,4 +868,6 @@ class RateLimitInfo:
         reset_delta = reset_at - self.reset_at
         self.reset_at = reset_at
         # TODO: remove
-        _logger.debug(f'Update {self.resource} rate limit: rem={self.remaining} reset={self.reset_at.strftime(TIME_FORMAT)} reset_delta={reset_delta.total_seconds()} recorded_delta={(self.reset_at - recorded_at).total_seconds()}')
+        _logger.debug(
+            f'Update {self.resource} rate limit: rem={self.remaining} reset={self.reset_at.strftime(TIME_FORMAT)} reset_delta={reset_delta.total_seconds()} recorded_delta={(self.reset_at - recorded_at).total_seconds()}'
+        )

pixeltable/exec/__init__.py

@@ -4,7 +4,7 @@ from .component_iteration_node import ComponentIterationNode
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext
 from .exec_node import ExecNode
+from .expr_eval import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
 from .row_update_node import RowUpdateNode
-from .sql_node import SqlLookupNode, SqlScanNode, SqlAggregationNode, SqlNode, SqlJoinNode
-from .expr_eval import ExprEvalNode
+from .sql_node import SqlAggregationNode, SqlJoinNode, SqlLookupNode, SqlNode, SqlScanNode

pixeltable/exec/aggregation_node.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import Any, Iterable, Iterator, Optional, cast, AsyncIterator
+from typing import Any, AsyncIterator, Iterable, Iterator, Optional, cast
 
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
@@ -13,12 +13,14 @@ from .exec_node import ExecNode
 
 _logger = logging.getLogger('pixeltable')
 
+
 class AggregationNode(ExecNode):
     """
     In-memory aggregation for UDAs.
 
     At the moment, this returns all results in a single DataRowBatch.
     """
+
     group_by: Optional[list[exprs.Expr]]
     input_exprs: list[exprs.Expr]
     agg_fn_eval_ctx: exprs.RowBuilder.EvalCtx
@@ -26,8 +28,13 @@ class AggregationNode(ExecNode):
     output_batch: DataRowBatch
 
     def __init__(
-            self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, group_by: Optional[list[exprs.Expr]],
-            agg_fn_calls: list[exprs.FunctionCall], input_exprs: Iterable[exprs.Expr], input: ExecNode
+        self,
+        tbl: catalog.TableVersion,
+        row_builder: exprs.RowBuilder,
+        group_by: Optional[list[exprs.Expr]],
+        agg_fn_calls: list[exprs.FunctionCall],
+        input_exprs: Iterable[exprs.Expr],
+        input: ExecNode,
     ):
         output_exprs: list[exprs.Expr] = [] if group_by is None else list(group_by)
         output_exprs.extend(agg_fn_calls)
@@ -86,4 +93,3 @@ class AggregationNode(ExecNode):
         self.output_batch.flush_imgs(None, self.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

@@ -9,7 +9,7 @@ import urllib.request
 from collections import deque
 from concurrent import futures
 from pathlib import Path
-from typing import Optional, Any, Iterator, AsyncIterator
+from typing import Any, AsyncIterator, Iterator, Optional
 from uuid import UUID
 
 import pixeltable.env as env
@@ -30,6 +30,7 @@ class CachePrefetchNode(ExecNode):
     TODO:
     - adapting the number of download threads at runtime to maximize throughput
     """
+
     BATCH_SIZE = 16
     NUM_EXECUTOR_THREADS = 16
 
@@ -59,8 +60,8 @@ class CachePrefetchNode(ExecNode):
         num_missing: int  # number of missing URLs in this row
 
     def __init__(
-            self, tbl_id: UUID, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode,
-            retain_input_order: bool = True):
+        self, tbl_id: UUID, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode, retain_input_order: bool = True
+    ):
         # input_/output_exprs=[]: we don't have anything to evaluate
         super().__init__(input.row_builder, [], [], input)
         self.retain_input_order = retain_input_order
@@ -241,6 +242,7 @@ class CachePrefetchNode(ExecNode):
             _logger.debug(f'Downloading {url} to {tmp_path}')
             if parsed.scheme == 's3':
                 from pixeltable.utils.s3 import get_client
+
                 with self.boto_client_lock:
                     if self.boto_client is None:
                         config = {

pixeltable/exec/component_iteration_node.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Iterator, Optional, AsyncIterator
+from typing import AsyncIterator, Iterator, Optional
 
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
@@ -14,6 +14,7 @@ class ComponentIterationNode(ExecNode):
 
     Returns row batches of OUTPUT_BATCH_SIZE size.
     """
+
     __OUTPUT_BATCH_SIZE = 1024
 
     def __init__(self, view: catalog.TableVersion, input: ExecNode):
@@ -25,8 +26,8 @@ class ComponentIterationNode(ExecNode):
         self.iterator_args = iterator_args[0]
         assert isinstance(self.iterator_args, exprs.InlineDict)
         self.iterator_args_ctx = self.row_builder.create_eval_ctx([self.iterator_args])
-        self.iterator_output_schema, self.unstored_column_names = (
-            self.view.iterator_cls.output_schema(**self.iterator_args.to_kwargs())
+        self.iterator_output_schema, self.unstored_column_names = self.view.iterator_cls.output_schema(
+            **self.iterator_args.to_kwargs()
         )
         self.iterator_output_fields = list(self.iterator_output_schema.keys())
         self.iterator_output_cols = {
@@ -34,7 +35,8 @@ class ComponentIterationNode(ExecNode):
         }
         # referenced iterator output fields
         self.refd_output_slot_idxs = {
-            e.col.name: e.slot_idx for e in self.row_builder.unique_exprs
+            e.col.name: e.slot_idx
+            for e in self.row_builder.unique_exprs
             if isinstance(e, exprs.ColumnRef) and e.col.name in self.iterator_output_fields
         }
 
@@ -79,8 +81,7 @@ class ComponentIterationNode(ExecNode):
         # verify and copy component_dict fields to their respective slots in output_row
         for field_name, field_val in component_dict.items():
             if field_name not in self.iterator_output_fields:
-                raise excs.Error(
-                    f'Invalid field name {field_name} in output of {self.view.iterator_cls.__name__}')
+                raise excs.Error(f'Invalid field name {field_name} in output of {self.view.iterator_cls.__name__}')
             if field_name not in self.refd_output_slot_idxs:
                 # we can ignore this
                 continue
@@ -90,5 +91,5 @@ class ComponentIterationNode(ExecNode):
         if len(component_dict) != len(self.iterator_output_fields):
             missing_fields = set(self.refd_output_slot_idxs.keys()) - set(component_dict.keys())
             raise excs.Error(
-                f'Invalid output of {self.view.iterator_cls.__name__}: '
-                f'missing fields {", ".join(missing_fields)}')
+                f'Invalid output of {self.view.iterator_cls.__name__}: missing fields {", ".join(missing_fields)}'
+            )

pixeltable/exec/data_row_batch.py

@@ -1,19 +1,21 @@
 from __future__ import annotations
-from typing import Iterator, Optional
+
 import logging
+from typing import Iterator, Optional
 
-import pixeltable.exprs as exprs
 import pixeltable.catalog as catalog
+import pixeltable.exprs as exprs
 from pixeltable.utils.media_store import MediaStore
 
-
 _logger = logging.getLogger('pixeltable')
 
+
 class DataRowBatch:
     """Set of DataRows, indexed by rowid.
 
     Contains the metadata needed to initialize DataRows.
     """
+
     tbl: Optional[catalog.TableVersion]
     row_builder: exprs.RowBuilder
     img_slot_idxs: list[int]
@@ -22,8 +24,11 @@ class DataRowBatch:
     rows: list[exprs.DataRow]
 
     def __init__(
-        self, tbl: Optional[catalog.TableVersion], row_builder: exprs.RowBuilder, num_rows: Optional[int] = None,
-        rows: Optional[list[exprs.DataRow]] = None
+        self,
+        tbl: Optional[catalog.TableVersion],
+        row_builder: exprs.RowBuilder,
+        num_rows: Optional[int] = None,
+        rows: Optional[list[exprs.DataRow]] = None,
     ):
         """
         Requires either num_rows or rows to be specified, but not both.
@@ -34,7 +39,8 @@ class DataRowBatch:
         self.img_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_image_type()]
         # non-image media slots
         self.media_slot_idxs = [
-            e.slot_idx for e in row_builder.unique_exprs
+            e.slot_idx
+            for e in row_builder.unique_exprs
             if e.col_type.is_media_type() and not e.col_type.is_image_type()
         ]
         self.array_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_array_type()]
@@ -44,14 +50,17 @@ class DataRowBatch:
             if num_rows is None:
                 num_rows = 0
             self.rows = [
-                exprs.DataRow(row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs)
+                exprs.DataRow(
+                    row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs
+                )
                 for _ in range(num_rows)
             ]
 
     def add_row(self, row: Optional[exprs.DataRow] = None) -> exprs.DataRow:
         if row is None:
             row = exprs.DataRow(
-                self.row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs)
+                self.row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs
+            )
         self.rows.append(row)
         return row
 
@@ -65,8 +74,10 @@ class DataRowBatch:
         return self.rows[index]
 
     def flush_imgs(
-            self, idx_range: Optional[slice] = None, stored_img_info: Optional[list[exprs.ColumnSlotIdx]] = None,
-            flushed_slot_idxs: Optional[list[int]] = None
+        self,
+        idx_range: Optional[slice] = None,
+        stored_img_info: Optional[list[exprs.ColumnSlotIdx]] = None,
+        flushed_slot_idxs: Optional[list[int]] = None,
     ) -> None:
         """Flushes images in the given range of rows."""
         assert self.tbl is not None

pixeltable/exec/exec_context.py

@@ -4,12 +4,19 @@ import sqlalchemy as sql
 
 import pixeltable.exprs as exprs
 
+
 class ExecContext:
     """Class for execution runtime constants"""
+
     def __init__(
-            self, row_builder: exprs.RowBuilder, *, show_pbar: bool = False, batch_size: int = 0,
-            pk_clause: Optional[list[sql.ClauseElement]] = None, num_computed_exprs: int = 0,
-            ignore_errors: bool = False
+        self,
+        row_builder: exprs.RowBuilder,
+        *,
+        show_pbar: bool = False,
+        batch_size: int = 0,
+        pk_clause: Optional[list[sql.ClauseElement]] = None,
+        num_computed_exprs: int = 0,
+        ignore_errors: bool = False,
     ):
         self.show_pbar = show_pbar
         self.batch_size = batch_size

pixeltable/exec/exec_node.py

@@ -4,16 +4,19 @@ import abc
 import asyncio
 import logging
 import sys
-from typing import Iterable, Iterator, Optional, TypeVar, AsyncIterator
+from typing import AsyncIterator, Iterable, Iterator, Optional, TypeVar
 
 import pixeltable.exprs as exprs
+
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext
 
 _logger = logging.getLogger('pixeltable')
 
+
 class ExecNode(abc.ABC):
     """Base class of all execution nodes"""
+
     output_exprs: Iterable[exprs.Expr]
     row_builder: exprs.RowBuilder
     input: Optional[ExecNode]
@@ -22,8 +25,12 @@ class ExecNode(abc.ABC):
     ctx: Optional[ExecContext]
 
     def __init__(
-            self, row_builder: exprs.RowBuilder, output_exprs: Iterable[exprs.Expr],
-            input_exprs: Iterable[exprs.Expr], input: Optional[ExecNode] = None):
+        self,
+        row_builder: exprs.RowBuilder,
+        output_exprs: Iterable[exprs.Expr],
+        input_exprs: Iterable[exprs.Expr],
+        input: Optional[ExecNode] = None,
+    ):
         self.output_exprs = output_exprs
         self.row_builder = row_builder
         self.input = input
@@ -31,8 +38,7 @@ class ExecNode(abc.ABC):
         output_slot_idxs = {e.slot_idx for e in output_exprs}
         output_dependencies = row_builder.get_dependencies(output_exprs, exclude=input_exprs)
         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
+            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
@@ -53,16 +59,20 @@ class ExecNode(abc.ABC):
         pass
 
     def __iter__(self) -> Iterator[DataRowBatch]:
+        running_loop: Optional[asyncio.AbstractEventLoop] = None
+        loop: asyncio.AbstractEventLoop
         try:
-            # check if we are already in an event loop (eg, Jupyter's); if so, patch it to allow nested event loops
-            _ = asyncio.get_event_loop()
+            # check if we are already in an event loop (eg, Jupyter's); if so, patch it to allow
+            # multiple run_until_complete()
+            running_loop = asyncio.get_running_loop()
             import nest_asyncio  # type: ignore
+
             nest_asyncio.apply()
+            loop = running_loop
+            _logger.debug(f'Patched running loop')
         except RuntimeError:
-            pass
-
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
 
         if 'pytest' in sys.modules:
             loop.set_debug(True)
@@ -75,7 +85,8 @@ class ExecNode(abc.ABC):
         except StopAsyncIteration:
             pass
         finally:
-            loop.close()
+            if loop != running_loop:
+                loop.close()
 
     def open(self) -> None:
         """Bottom-up initialization of nodes for execution. Must be called before __next__."""

pixeltable/exec/expr_eval/evaluators.py

@@ -5,10 +5,10 @@ import datetime
 import itertools
 import logging
 import sys
-from typing import Iterator, Any, Optional, Callable, cast
+from typing import Any, Callable, Iterator, Optional, cast
+
+from pixeltable import exprs, func
 
-from pixeltable import exprs
-from pixeltable import func
 from .globals import Dispatcher, Evaluator, FnCallArgs
 
 _logger = logging.getLogger('pixeltable')
@@ -23,6 +23,7 @@ class DefaultExprEvaluator(Evaluator):
     TODO:
     - parallelize via Ray
     """
+
     e: exprs.Expr
 
     def __init__(self, e: exprs.Expr, dispatcher: Dispatcher):
@@ -32,8 +33,7 @@ class DefaultExprEvaluator(Evaluator):
     def schedule(self, rows: list[exprs.DataRow], slot_idx: int) -> None:
         assert self.e.slot_idx >= 0
         task = asyncio.create_task(self.eval(rows))
-        self.dispatcher.tasks.add(task)
-        task.add_done_callback(self.dispatcher.done_cb)
+        self.dispatcher.register_task(task)
 
     async def eval(self, rows: list[exprs.DataRow]) -> None:
         rows_with_excs: set[int] = set()  # records idxs into rows
@@ -61,6 +61,7 @@ class FnCallEvaluator(Evaluator):
     TODO:
     - adaptive batching: finding the optimal batch size based on observed execution times
     """
+
     fn_call: exprs.FunctionCall
     fn: func.CallableFunction
     scalar_py_fn: Optional[Callable]  # only set for non-batching CallableFunctions
@@ -74,7 +75,7 @@ class FnCallEvaluator(Evaluator):
         self.fn_call = fn_call
         self.fn = cast(func.CallableFunction, fn_call.fn)
         if isinstance(self.fn, func.CallableFunction) and self.fn.is_batched:
-            self.call_args_queue =  asyncio.Queue[FnCallArgs]()
+            self.call_args_queue = asyncio.Queue[FnCallArgs]()
             # we're not supplying sample arguments there, they're ignored anyway
             self.batch_size = self.fn.get_batch_size()
             self.scalar_py_fn = None
@@ -134,8 +135,7 @@ class FnCallEvaluator(Evaluator):
                     scheduler.submit(batched_call_args)
                 else:
                     task = asyncio.create_task(self.eval_batch(batched_call_args))
-                    self.dispatcher.tasks.add(task)
-                    task.add_done_callback(self.dispatcher.done_cb)
+                    self.dispatcher.register_task(task)
 
         elif self.fn.is_async:
             if self.fn_call.resource_pool is not None:
@@ -147,14 +147,12 @@ class FnCallEvaluator(Evaluator):
                 # create one task per call
                 for item in rows_call_args:
                     task = asyncio.create_task(self.eval_async(item))
-                    self.dispatcher.tasks.add(task)
-                    task.add_done_callback(self.dispatcher.done_cb)
+                    self.dispatcher.register_task(task)
 
         else:
             # create a single task for all rows
             task = asyncio.create_task(self.eval(rows_call_args))
-            self.dispatcher.tasks.add(task)
-            task.add_done_callback(self.dispatcher.done_cb)
+            self.dispatcher.register_task(task)
 
     def _queued_call_args_iter(self) -> Iterator[FnCallArgs]:
         while not self.call_args_queue.empty():
@@ -171,14 +169,16 @@ class FnCallEvaluator(Evaluator):
             for k in item.kwargs.keys():
                 batch_kwargs[k][i] = item.kwargs[k]
         return FnCallArgs(
-            self.fn_call, [item.row for item in call_args], batch_args=batch_args, batch_kwargs=batch_kwargs)
+            self.fn_call, [item.row for item in call_args], batch_args=batch_args, batch_kwargs=batch_kwargs
+        )
 
     async def eval_batch(self, batched_call_args: FnCallArgs) -> None:
         result_batch: list[Any]
         try:
             if self.fn.is_async:
                 result_batch = await self.fn.aexec_batch(
-                    *batched_call_args.batch_args, **batched_call_args.batch_kwargs)
+                    *batched_call_args.batch_args, **batched_call_args.batch_kwargs
+                )
             else:
                 # check for cancellation before starting something potentially long-running
                 if asyncio.current_task().cancelled() or self.dispatcher.exc_event.is_set():
@@ -209,6 +209,7 @@ class FnCallEvaluator(Evaluator):
             self.dispatcher.dispatch([call_args.row])
         except Exception as exc:
             import anthropic
+
             if isinstance(exc, anthropic.RateLimitError):
                 _logger.debug(f'RateLimitError: {exc}')
             _, _, exc_tb = sys.exc_info()
@@ -232,7 +233,8 @@ class FnCallEvaluator(Evaluator):
                 rows_with_excs.add(idx)
                 self.dispatcher.dispatch_exc(item.rows, self.fn_call.slot_idx, exc_tb)
         self.dispatcher.dispatch(
-            [call_args_batch[i].row for i in range(len(call_args_batch)) if i not in rows_with_excs])
+            [call_args_batch[i].row for i in range(len(call_args_batch)) if i not in rows_with_excs]
+        )
 
     def _close(self) -> None:
         """Create a task for the incomplete batch of queued FnCallArgs, if any"""
@@ -241,5 +243,4 @@ class FnCallEvaluator(Evaluator):
             return
         batched_call_args = self._create_batch_call_args(list(self._queued_call_args_iter()))
         task = asyncio.create_task(self.eval_batch(batched_call_args))
-        self.dispatcher.tasks.add(task)
-        task.add_done_callback(self.dispatcher.done_cb)
+        self.dispatcher.register_task(task)