pixeltable 0.1.0__py3-none-any.whl → 0.2.4__py3-none-any.whl

pixeltable/env.py

@@ -1,86 +1,400 @@
-from typing import Optional
+from __future__ import annotations
+
+import datetime
+import glob
+import http.server
+import importlib
+import importlib.util
+import logging
+import os
+import socketserver
+import sys
+import threading
+import typing
+import uuid
 from pathlib import Path
+from typing import Callable, Optional, Dict, Any, List
+
+import pgserver
 import sqlalchemy as sql
+import yaml
 from sqlalchemy_utils.functions import database_exists, create_database, drop_database
 
+import pixeltable.exceptions as excs
+from pixeltable import metadata
+
 
 class Env:
     """
     Store for runtime globals.
     """
-    _instance: Optional['Env'] = None
+    _instance: Optional[Env] = None
+    _log_fmt_str = '%(asctime)s %(levelname)s %(name)s %(filename)s:%(lineno)d: %(message)s'
 
     @classmethod
-    def get(cls) -> 'Env':
+    def get(cls) -> Env:
         if cls._instance is None:
             cls._instance = Env()
         return cls._instance
 
     def __init__(self):
         self._home: Optional[Path] = None
-        self._db_path: Optional[Path] = None
-        self._img_dir: Optional[Path] = None
-        self._nnidx_dir: Optional[Path] = None
-        self._tmp_video_dir: Optional[Path] = None
+        self._media_dir: Optional[Path] = None  # computed media files
+        self._file_cache_dir: Optional[Path] = None  # cached media files with external URL
+        self._dataset_cache_dir: Optional[Path] = None  # cached datasets (eg, pytorch or COCO)
+        self._log_dir: Optional[Path] = None  # log files
+        self._tmp_dir: Optional[Path] = None  # any tmp files
         self._sa_engine: Optional[sql.engine.base.Engine] = None
+        self._pgdata_dir : Optional[Path] = None
         self._db_name: Optional[str] = None
+        self._db_server: Optional[pgserver.PostgresServer] = None
+        self._db_url: Optional[str] = None
+
+        # info about installed packages that are utilized by some parts of the code;
+        # package name -> version; version == []: package is installed, but we haven't determined the version yet
+        self._installed_packages: Dict[str, Optional[List[int]]] = {}
+        self._nos_client: Optional[Any] = None
+        self._spacy_nlp: Optional[Any] = None  # spacy.Language
+        self._httpd: Optional[socketserver.TCPServer] = None
+        self._http_address: Optional[str] = None
+
+        self._registered_clients: dict[str, Any] = {}
+
+        # logging-related state
+        self._logger = logging.getLogger('pixeltable')
+        self._logger.setLevel(logging.DEBUG)  # allow everything to pass, we filter in _log_filter()
+        self._logger.propagate = False
+        self._logger.addFilter(self._log_filter)
+        self._default_log_level = logging.INFO
+        self._logfilename: Optional[str] = None
+        self._log_to_stdout = False
+        self._module_log_level: Dict[str, int] = {}  # module name -> log level
+
+        # config
+        self._config_file: Optional[Path] = None
+        self._config: Optional[Dict[str, Any]] = None
+
+        # create logging handler to also log to stdout
+        self._stdout_handler = logging.StreamHandler(stream=sys.stdout)
+        self._stdout_handler.setFormatter(logging.Formatter(self._log_fmt_str))
+        self._initialized = False
+
+    @property
+    def config(self):
+        return self._config
+
+    @property
+    def db_url(self) -> str:
+        assert self._db_url is not None
+        return self._db_url
+
+    @property
+    def http_address(self) -> str:
+        assert self._http_address is not None
+        return self._http_address
+
+    def print_log_config(self) -> None:
+        print(f'logging to {self._logfilename}')
+        print(f'{"" if self._log_to_stdout else "not "}logging to stdout')
+        print(f'default log level: {logging.getLevelName(self._default_log_level)}')
+        print(
+            f'module log levels: '
+            f'{",".join([name + ":" + logging.getLevelName(val) for name, val in self._module_log_level.items()])}')
+
+    def log_to_stdout(self, enable: bool = True) -> None:
+        self._log_to_stdout = enable
+        if enable:
+            self._logger.addHandler(self._stdout_handler)
+        else:
+            self._logger.removeHandler(self._stdout_handler)
+
+    def set_log_level(self, level: int) -> None:
+        self._default_log_level = level
+
+    def set_module_log_level(self, module: str, level: Optional[int]) -> None:
+        if level is None:
+            self._module_log_level.pop(module, None)
+        else:
+            self._module_log_level[module] = level
+
+    def is_installed_package(self, package_name: str) -> bool:
+        return self._installed_packages[package_name] is not None
+
+    def _log_filter(self, record: logging.LogRecord) -> bool:
+        if record.name == 'pixeltable':
+            # accept log messages from a configured pixeltable module (at any level of the module hierarchy)
+            path_parts = list(Path(record.pathname).parts)
+            path_parts.reverse()
+            max_idx = path_parts.index('pixeltable')
+            for module_name in path_parts[:max_idx]:
+                if module_name in self._module_log_level and record.levelno >= self._module_log_level[module_name]:
+                    return True
+        if record.levelno >= self._default_log_level:
+            return True
+        else:
+            return False
+
+    def set_up(self, echo: bool = False, reinit_db: bool = False) -> None:
+        if self._initialized:
+            return
+
+        self._initialized = True
+        home = Path(os.environ.get('PIXELTABLE_HOME', str(Path.home() / '.pixeltable')))
+        assert self._home is None or self._home == home
+        self._home = home
+        self._config_file = Path(os.environ.get('PIXELTABLE_CONFIG', str(self._home / 'config.yaml')))
+        self._media_dir = self._home / 'media'
+        self._file_cache_dir = self._home / 'file_cache'
+        self._dataset_cache_dir = self._home / 'dataset_cache'
+        self._log_dir = self._home / 'logs'
+        self._tmp_dir = self._home / 'tmp'
+
+        # Read in the config
+        if os.path.isfile(self._config_file):
+            with open(self._config_file, 'r') as stream:
+                try:
+                    self._config = yaml.safe_load(stream)
+                except yaml.YAMLError as exc:
+                    self._logger.error(f'Could not read config file: {self._config_file}')
+                    self._config = {}
+        else:
+            self._config = {}
 
-    def set_up(
-        self, home_parent: Optional[Path] = Path.home(), db_name: str = 'pixeltable', echo: bool = False
-    ) -> None:
-        self.set_home(home_parent / '.pixeltable')
         if self._home.exists() and not self._home.is_dir():
             raise RuntimeError(f'{self._home} is not a directory')
 
-        self._db_name = db_name
-        db_url = f'postgresql:///{self._db_name}'
-
         if not self._home.exists():
-            print(f'creating {self._home}')
+            # we don't have our logger set up yet, so print to stdout
+            print(f'Creating a Pixeltable instance at: {self._home}')
             self._home.mkdir()
-            self._img_dir.mkdir()
-            self._nnidx_dir.mkdir()
-            self._tmp_video_dir.mkdir()
-            self.tear_down()
-            if not database_exists(db_url):
-                create_database(db_url)
-            self._sa_engine = sql.create_engine(db_url, echo=echo, future=True)
-            from pixeltable import store
-            store.Base.metadata.create_all(self._sa_engine)
+            # TODO (aaron-siegel) This is the existing behavior, but it seems scary. If something happens to
+            # self._home, it will cause the DB to be destroyed even if pgdata is in an alternate location.
+            # PROPOSAL: require `reinit_db` to be set explicitly to destroy the DB.
+            reinit_db = True
+
+        if not self._media_dir.exists():
+            self._media_dir.mkdir()
+        if not self._file_cache_dir.exists():
+            self._file_cache_dir.mkdir()
+        if not self._dataset_cache_dir.exists():
+            self._dataset_cache_dir.mkdir()
+        if not self._log_dir.exists():
+            self._log_dir.mkdir()
+        if not self._tmp_dir.exists():
+            self._tmp_dir.mkdir()
+
+        # 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')
+        fh.setFormatter(logging.Formatter(self._log_fmt_str))
+        self._logger.addHandler(fh)
+        sql_logger = logging.getLogger('sqlalchemy.engine')
+        sql_logger.setLevel(logging.INFO)
+        sql_logger.addHandler(fh)
+        sql_logger.propagate = False
+
+        # empty tmp dir
+        for path in glob.glob(f'{self._tmp_dir}/*'):
+            os.remove(path)
+
+        self._db_name = os.environ.get('PIXELTABLE_DB', 'pixeltable')
+        self._pgdata_dir = Path(os.environ.get('PIXELTABLE_PGDATA', str(self._home / 'pgdata')))
+
+        # in pgserver.get_server(): cleanup_mode=None will leave db on for debugging purposes
+        self._db_server = pgserver.get_server(self._pgdata_dir, cleanup_mode=None)
+        self._db_url = self._db_server.get_uri(database=self._db_name)
+
+        if reinit_db:
+            if database_exists(self.db_url):
+                drop_database(self.db_url)
+
+        if not database_exists(self.db_url):
+            self._logger.info(f'creating database at {self.db_url}')
+            create_database(self.db_url)
+            self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True)
+            from pixeltable.metadata import schema
+            schema.Base.metadata.create_all(self._sa_engine)
+            metadata.create_system_info(self._sa_engine)
+            # enable pgvector
+            with self._sa_engine.begin() as conn:
+                conn.execute(sql.text('CREATE EXTENSION vector'))
         else:
+            self._logger.info(f'found database {self.db_url}')
             if self._sa_engine is None:
-                self._sa_engine = sql.create_engine(db_url, echo=echo, future=True)
+                self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True)
+
+        print(f'Connected to Pixeltable database at: {self.db_url}')
 
-    def tear_down(self) -> None:
-        db_url = f'postgresql:///{self._db_name}'
-        if database_exists(db_url):
-            drop_database(db_url)
+        # we now have a home directory and db; start other services
+        self._set_up_runtime()
+        self.log_to_stdout(False)
 
-    def set_home(self, home: Path) -> None:
-        if self._home is not None:
+    def upgrade_metadata(self) -> None:
+        metadata.upgrade_md(self._sa_engine)
+
+    def _create_nos_client(self) -> None:
+        import nos
+        self._logger.info('connecting to NOS')
+        nos.init(logging_level=logging.DEBUG)
+        self._nos_client = nos.client.InferenceClient()
+        self._logger.info('waiting for NOS')
+        self._nos_client.WaitForServer()
+
+        # now that we have a client, we can create the module
+        import importlib
+        try:
+            importlib.import_module('pixeltable.functions.nos')
+            # it's already been created
             return
-        self._home = home
-        self._db_path = self._home / 'db.sqlite3'
-        self._img_dir = self._home / 'images'
-        self._nnidx_dir = self._home / 'nnidxs'
-        self._tmp_video_dir = self._home / 'tmp_videos'
+        except ImportError:
+            pass
+        from pixeltable.functions.util import create_nos_modules
+        _ = create_nos_modules()
+
+    def get_client(self, name: str, init: Callable, environ: Optional[str] = None) -> Any:
+        """
+        Gets the client with the specified name, using `init` to construct one if necessary.
+
+        - name: The name of the client
+        - init: A `Callable` with signature `fn(api_key: str) -> Any` that constructs a client object
+        - environ: The name of the environment variable to use for the API key, if no API key is found in config
+            (defaults to f'{name.upper()}_API_KEY')
+        """
+        if name in self._registered_clients:
+            return self._registered_clients[name]
+
+        if environ is None:
+            environ = f'{name.upper()}_API_KEY'
+
+        if name in self._config and 'api_key' in self._config[name]:
+            api_key = self._config[name]['api_key']
+        else:
+            api_key = os.environ.get(environ)
+        if api_key is None or api_key == '':
+            raise excs.Error(f'`{name}` client not initialized (no API key configured).')
+
+        client = init(api_key)
+        self._registered_clients[name] = client
+        self._logger.info(f'Initialized `{name}` client.')
+        return client
+
+    def _start_web_server(self) -> None:
+        """
+        The http server root is the file system root.
+        eg: /home/media/foo.mp4 is located at http://127.0.0.1:{port}/home/media/foo.mp4
+        This arrangement enables serving media hosted within _home,
+        as well as external media inserted into pixeltable or produced by pixeltable.
+        The port is chosen dynamically to prevent conflicts.
+        """
+        # Port 0 means OS picks one for us.
+        address = ("127.0.0.1", 0)
+        class FixedRootHandler(http.server.SimpleHTTPRequestHandler):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, directory='/', **kwargs)
+        self._httpd = socketserver.TCPServer(address, FixedRootHandler)
+        port = self._httpd.server_address[1]
+        self._http_address = f'http://127.0.0.1:{port}'
+
+        def run_server():
+            logging.log(logging.INFO, f'running web server at {self._http_address}')
+            self._httpd.serve_forever()
+
+        # Run the server in a separate thread
+        thread = threading.Thread(target=run_server, daemon=True)
+        thread.start()
+
+    def _set_up_runtime(self) -> None:
+        """Check for and start runtime services"""
+        self._start_web_server()
+        self._check_installed_packages()
+
+    def _check_installed_packages(self) -> None:
+        def check(package: str) -> None:
+            if importlib.util.find_spec(package) is not None:
+                self._installed_packages[package] = []
+            else:
+                self._installed_packages[package] = None
+
+        check('datasets')
+        check('torch')
+        check('torchvision')
+        check('transformers')
+        check('sentence_transformers')
+        check('boto3')
+        check('pyarrow')
+        check('spacy')  # TODO: deal with en-core-web-sm
+        if self.is_installed_package('spacy'):
+            import spacy
+            self._spacy_nlp = spacy.load('en_core_web_sm')
+        check('tiktoken')
+        check('openai')
+        check('together')
+        check('fireworks')
+        check('nos')
+        if self.is_installed_package('nos'):
+            self._create_nos_client()
+
+    def require_package(self, package: str, min_version: Optional[List[int]] = None) -> None:
+        assert package in self._installed_packages
+        if self._installed_packages[package] is None:
+            raise excs.Error(f'Package {package} is not installed')
+        if min_version is None:
+            return
+
+        # check whether we have a version >= the required one
+        if self._installed_packages[package] == []:
+            m = importlib.import_module(package)
+            module_version = [int(x) for x in m.__version__.split('.')]
+            self._installed_packages[package] = module_version
+        installed_version = self._installed_packages[package]
+        if len(min_version) < len(installed_version):
+            normalized_min_version = min_version + [0] * (len(installed_version) - len(min_version))
+        if any([a < b for a, b in zip(installed_version, normalized_min_version)]):
+            raise excs.Error((
+                f'The installed version of package {package} is {".".join([str[v] for v in installed_version])}, '
+                f'but version  >={".".join([str[v] for v in min_version])} is required'))
+
+    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}'
 
     @property
-    def img_dir(self) -> Path:
-        assert self._img_dir is not None
-        return self._img_dir
+    def home(self) -> Path:
+        assert self._home is not None
+        return self._home
 
     @property
-    def nnidx_dir(self) -> Path:
-        assert self._nnidx_dir is not None
-        return self._nnidx_dir
+    def media_dir(self) -> Path:
+        assert self._media_dir is not None
+        return self._media_dir
 
     @property
-    def tmp_video_dir(self) -> Path:
-        assert self._tmp_video_dir is not None
-        return self._tmp_video_dir
+    def file_cache_dir(self) -> Path:
+        assert self._file_cache_dir is not None
+        return self._file_cache_dir
+
+    @property
+    def dataset_cache_dir(self) -> Path:
+        assert self._dataset_cache_dir is not None
+        return self._dataset_cache_dir
+
+    @property
+    def tmp_dir(self) -> Path:
+        assert self._tmp_dir is not None
+        return self._tmp_dir
 
     @property
     def engine(self) -> sql.engine.base.Engine:
         assert self._sa_engine is not None
         return self._sa_engine
+
+    @property
+    def nos_client(self) -> Any:
+        return self._nos_client
+
+    @property
+    def spacy_nlp(self) -> Any:
+        assert self._spacy_nlp is not None
+        return self._spacy_nlp

pixeltable/exceptions.py

@@ -1,26 +1,17 @@
-class Error(Exception):
-    pass
-
-
-class DuplicateNameError(Exception):
-    pass
-
-
-class UnknownEntityError(Exception):
-    pass
+from typing import List, Any
+from types import TracebackType
+from dataclasses import dataclass
 
 
-class BadFormatError(Exception):
-    pass
-
-
-class DirectoryNotEmptyError(Exception):
-    pass
-
-
-class InsertError(Exception):
+class Error(Exception):
     pass
 
 
-class OperationalError(Exception):
-    pass
+@dataclass
+class ExprEvalError(Exception):
+    expr: Any  # exprs.Expr, but we're not importing pixeltable.exprs to avoid circular imports
+    expr_msg: str
+    exc: Exception
+    exc_tb: TracebackType
+    input_vals: List[Any]
+    row_num: int

pixeltable/exec/__init__.py

@@ -0,0 +1,9 @@
+from .aggregation_node import AggregationNode
+from .cache_prefetch_node import CachePrefetchNode
+from .component_iteration_node import ComponentIterationNode
+from .exec_context import ExecContext
+from .exec_node import ExecNode
+from .expr_eval_node import ExprEvalNode
+from .in_memory_data_node import InMemoryDataNode
+from .sql_scan_node import SqlScanNode
+from .media_validation_node import MediaValidationNode

pixeltable/exec/aggregation_node.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import logging
+import sys
+from typing import List, Optional, Any
+
+import pixeltable.catalog as catalog
+import pixeltable.exceptions as excs
+import pixeltable.exprs as exprs
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+
+_logger = logging.getLogger('pixeltable')
+
+class AggregationNode(ExecNode):
+    def __init__(
+            self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, group_by: List[exprs.Expr],
+            agg_fn_calls: List[exprs.FunctionCall], input_exprs: List[exprs.Expr], input: ExecNode
+    ):
+        super().__init__(row_builder, group_by + agg_fn_calls, input_exprs, input)
+        self.input = input
+        self.group_by = group_by
+        self.input_exprs = input_exprs
+        self.agg_fn_calls = agg_fn_calls
+        self.agg_fn_eval_ctx = row_builder.create_eval_ctx(agg_fn_calls, exclude=input_exprs)
+        self.output_batch = DataRowBatch(tbl, row_builder, 0)
+
+    def _reset_agg_state(self, row_num: int) -> None:
+        for fn_call in self.agg_fn_calls:
+            try:
+                fn_call.reset_agg()
+            except Exception as e:
+                _, _, exc_tb = sys.exc_info()
+                expr_msg = f'init() function of the aggregate {fn_call}'
+                raise excs.ExprEvalError(fn_call, expr_msg, e, exc_tb, [], row_num)
+
+    def _update_agg_state(self, row: exprs.DataRow, row_num: int) -> None:
+        for fn_call in self.agg_fn_calls:
+            try:
+                fn_call.update(row)
+            except Exception as e:
+                _, _, exc_tb = sys.exc_info()
+                expr_msg = f'update() function of the aggregate {fn_call}'
+                input_vals = [row[d.slot_idx] for d in fn_call.dependencies()]
+                raise excs.ExprEvalError(fn_call, expr_msg, e, exc_tb, input_vals, row_num)
+
+    def __next__(self) -> DataRowBatch:
+        if self.output_batch is None:
+            raise StopIteration
+
+        prev_row: Optional[exprs.DataRow] = None
+        current_group: Optional[List[Any]] = None  # the values of the group-by exprs
+        num_input_rows = 0
+        for row_batch in self.input:
+            num_input_rows += len(row_batch)
+            for row in row_batch:
+                group = [row[e.slot_idx] for e in self.group_by]
+                if current_group is None:
+                    current_group = group
+                    self._reset_agg_state(0)
+                if group != current_group:
+                    # we're entering a new group, emit a row for the previous one
+                    self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
+                    self.output_batch.add_row(prev_row)
+                    current_group = group
+                    self._reset_agg_state(0)
+                self._update_agg_state(row, 0)
+                prev_row = row
+        # emit the last group
+        self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
+        self.output_batch.add_row(prev_row)
+
+        result = self.output_batch
+        result.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
+        self.output_batch = None
+        _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(result.rows)} rows')
+        return result
+