Merge branch 'master' of https://github.com/datajoint/datajoint-python into populate-kwargs

Author: A-Baji

CHANGELOG.md

@@ -5,14 +5,17 @@
 * Bugfix - Fix Python 3.10 compatibility (#983) PR #972
 * Bugfix - Allow renaming non-conforming attributes in proj (#982) PR #972
 * Add - Expose proxy feature for S3 external stores (#961) PR #962
+* Add - implement multiprocessing in populate (#695) PR #704, #969
 * Bugfix - Dependencies not properly loaded on populate. (#902) PR #919
 * Bugfix - Replace use of numpy aliases of built-in types with built-in type. (#938) PR #939
+* Bugfix - Deletes and drops must include the master of each part. (#151, #374) PR #957
 * Bugfix - `ExternalTable.delete` should not remove row on error (#953) PR #956
 * Bugfix - Fix error handling of remove_object function in `s3.py` (#952) PR #955
 * Bugfix - Fix regression issue with `DISTINCT` clause and `GROUP_BY` (#914) PR #963
 * Bugfix - Fix sql code generation to comply with sql mode `ONLY_FULL_GROUP_BY` (#916) PR #965
 * Bugfix - Fix count for left-joined `QueryExpressions` (#951) PR #966
 * Bugfix - Fix assertion error when performing a union into a join (#930) PR #967
+* Update `~jobs.error_stack` from blob to mediumblob to allow error stacks >64kB in jobs (#984) PR #986
 
 ### 0.13.2 -- May 7, 2021
 * Update `setuptools_certificate` dependency to new name `otumat`
@@ -133,7 +136,7 @@
 * Fix #628 - incompatibility with pyparsing 2.4.1
 
 ### 0.11.1 -- Nov 15, 2018
-* Fix ordering of attributes in proj (#483 and #516)
+* Fix ordering of attributes in proj (#483, #516)
 * Prohibit direct insert into auto-populated tables (#511)
 
 ### 0.11.0 -- Oct 25, 2018
@@ -246,7 +249,7 @@ Documentation and tutorials available at https://docs.datajoint.io and https://t
 * ERD() no longer text the context argument.
 * ERD.draw() now takes an optional context argument.  By default uses the caller's locals.
 
-### 0.3.2.   
+### 0.3.2.
 * Fixed issue #223:  `insert` can insert relations without fetching.
 * ERD() now takes the `context` argument, which specifies in which context to look for classes. The default is taken from the argument (schema or relation).
 * ERD.draw() no longer has the `prefix` argument: class names are shown as found in the context.

datajoint/autopopulate.py

@@ -8,12 +8,37 @@
 from .expression import QueryExpression, AndList
 from .errors import DataJointError, LostConnectionError
 import signal
+import multiprocessing as mp
 
 # noinspection PyExceptionInherit,PyCallingNonCallable
 
 logger = logging.getLogger(__name__)
 
 
+# --- helper functions for multiprocessing --
+
+def _initialize_populate(table, jobs, populate_kwargs):
+    """
+    Initialize the process for mulitprocessing.
+    Saves the unpickled copy of the table to the current process and reconnects.
+    """
+    process = mp.current_process()
+    process.table = table
+    process.jobs = jobs
+    process.populate_kwargs = populate_kwargs
+    table.connection.connect()  # reconnect
+
+
+def _call_populate1(key):
+    """
+    Call current process' table._populate1()
+    :key - a dict specifying job to compute
+    :return: key, error if error, otherwise None
+    """
+    process = mp.current_process()
+    return process.table._populate1(key, process.jobs, **process.populate_kwargs)
+
+
 class AutoPopulate:
     """
     AutoPopulate is a mixin class that adds the method populate() to a Relation class.
@@ -26,18 +51,20 @@ class AutoPopulate:
     @property
     def key_source(self):
         """
-        :return: the relation whose primary key values are passed, sequentially, to the
-                ``make`` method when populate() is called.
-                The default value is the join of the parent relations.
-                Users may override to change the granularity or the scope of populate() calls.
+        :return: the query expression that yields primary key values to be passed,
+        sequentially, to the ``make`` method when populate() is called.
+        The default value is the join of the parent tables references from the primary key.
+        Subclasses may override they key_source to change the scope or the granularity
+        of the make calls.
         """
         def _rename_attributes(table, props):
             return (table.proj(
                 **{attr: ref for attr, ref in props['attr_map'].items() if attr != ref})
-                if props['aliased'] else table)
+                if props['aliased'] else table.proj())
 
         if self._key_source is None:
-            parents = self.target.parents(primary=True, as_objects=True, foreign_key_info=True)
+            parents = self.target.parents(
+                primary=True, as_objects=True, foreign_key_info=True)
             if not parents:
                 raise DataJointError('A table must have dependencies '
                                      'from its primary key for auto-populate to work')
@@ -48,17 +75,19 @@ def _rename_attributes(table, props):
 
     def make(self, key):
         """
-        Derived classes must implement method `make` that fetches data from tables that are
-        above them in the dependency hierarchy, restricting by the given key, computes dependent
-        attributes, and inserts the new tuples into self.
+        Derived classes must implement method `make` that fetches data from tables
+        above them in the dependency hierarchy, restricting by the given key,
+        computes secondary attributes, and inserts the new tuples into self.
         """
-        raise NotImplementedError('Subclasses of AutoPopulate must implement the method `make`')
+        raise NotImplementedError(
+            'Subclasses of AutoPopulate must implement the method `make`')
 
     @property
     def target(self):
         """
         :return: table to be populated.
-        In the typical case, dj.AutoPopulate is mixed into a dj.Table class by inheritance and the target is self.
+        In the typical case, dj.AutoPopulate is mixed into a dj.Table class by
+        inheritance and the target is self.
         """
         return self
 
@@ -85,41 +114,49 @@ def _jobs_to_do(self, restrictions):
 
         if not isinstance(todo, QueryExpression):
             raise DataJointError('Invalid key_source value')
-        # check if target lacks any attributes from the primary key of key_source
+
         try:
+            # check if target lacks any attributes from the primary key of key_source
             raise DataJointError(
-                'The populate target lacks attribute %s from the primary key of key_source' % next(
-                    name for name in todo.heading.primary_key if name not in self.target.heading))
+                'The populate target lacks attribute %s '
+                'from the primary key of key_source' % next(
+                    name for name in todo.heading.primary_key
+                    if name not in self.target.heading))
         except StopIteration:
             pass
         return (todo & AndList(restrictions)).proj()
 
     def populate(self, *restrictions, suppress_errors=False, return_exception_objects=False,
                  reserve_jobs=False, order="original", limit=None, max_calls=None,
-                 display_progress=False, make_kwargs=None):
+                 display_progress=False, processes=1, make_kwargs=None):
         """
-        rel.populate() calls rel.make(key) for every primary key in self.key_source
-        for which there is not already a tuple in rel.
-        :param restrictions: a list of restrictions each restrict (rel.key_source - target.proj())
+        table.populate() calls table.make(key) for every primary key in self.key_source
+        for which there is not already a tuple in table.
+        :param restrictions: a list of restrictions each restrict
+            (table.key_source - target.proj())
         :param suppress_errors: if True, do not terminate execution.
         :param return_exception_objects: return error objects instead of just error messages
-        :param reserve_jobs: if true, reserves job to populate in asynchronous fashion
+        :param reserve_jobs: if True, reserve jobs to populate in asynchronous fashion
         :param order: "original"|"reverse"|"random"  - the order of execution
+        :param limit: if not None, check at most this many keys
+        :param max_calls: if not None, populate at most this many keys
         :param display_progress: if True, report progress_bar
         :param limit: if not None, checks at most that many keys
         :param max_calls: if not None, populates at max that many keys
-        :param make_kwargs: optional dict containing keyword arguments that will be passed down to each make() call
+        :param processes: number of processes to use. When set to a large number, then
+            uses as many as CPU cores
+        :param make_kwargs: optional dict containing keyword arguments that will be 
+            passed down to each make() call
         """
         if self.connection.in_transaction:
             raise DataJointError('Populate cannot be called during a transaction.')
 
         valid_order = ['original', 'reverse', 'random']
         if order not in valid_order:
             raise DataJointError('The order argument must be one of %s' % str(valid_order))
-        error_list = [] if suppress_errors else None
         jobs = self.connection.schemas[self.target.database].jobs if reserve_jobs else None
 
-        # define and setup signal handler for SIGTERM
+        # define and set up signal handler for SIGTERM:
         if reserve_jobs:
             def handler(signum, frame):
                 logger.info('Populate terminated by SIGTERM')
@@ -132,60 +169,99 @@ def handler(signum, frame):
         elif order == "random":
             random.shuffle(keys)
 
-        call_count = 0
         logger.info('Found %d keys to populate' % len(keys))
 
-        make = self._make_tuples if hasattr(self, '_make_tuples') else self.make
+        keys = keys[:max_calls]
+        nkeys = len(keys)
 
-        for key in (tqdm(keys, desc=self.__class__.__name__) if display_progress else keys):
-            if max_calls is not None and call_count >= max_calls:
-                break
-            if not reserve_jobs or jobs.reserve(self.target.table_name, self._job_key(key)):
-                self.connection.start_transaction()
-                if key in self.target:  # already populated
-                    self.connection.cancel_transaction()
-                    if reserve_jobs:
-                        jobs.complete(self.target.table_name, self._job_key(key))
+        if processes > 1:
+            processes = min(processes, nkeys, mp.cpu_count())
+
+        error_list = []
+        populate_kwargs = dict(
+            suppress_errors=suppress_errors,
+            return_exception_objects=return_exception_objects)
+
+        if processes == 1:
+            for key in tqdm(keys, desc=self.__class__.__name__) if display_progress else keys:
+                error = self._populate1(key, jobs, **populate_kwargs)
+                if error is not None:
+                    error_list.append(error)
+        else:
+            # spawn multiple processes
+            self.connection.close()  # disconnect parent process from MySQL server
+            del self.connection._conn.ctx  # SSLContext is not pickleable
+            with mp.Pool(processes, _initialize_populate, (self, populate_kwargs)) as pool:
+                if display_progress:
+                    with tqdm(desc="Processes: ", total=nkeys) as pbar:
+                        for error in pool.imap(_call_populate1, keys, chunksize=1):
+                            if error is not None:
+                                error_list.append(error)
+                            pbar.update()
                 else:
-                    logger.info('Populating: ' + str(key))
-                    call_count += 1
-                    self.__class__._allow_insert = True
-                    try:
-                        make(dict(key), **(make_kwargs or {}))
-                    except (KeyboardInterrupt, SystemExit, Exception) as error:
-                        try:
-                            self.connection.cancel_transaction()
-                        except LostConnectionError:
-                            pass
-                        error_message = '{exception}{msg}'.format(
-                            exception=error.__class__.__name__,
-                            msg=': ' + str(error) if str(error) else '')
-                        if reserve_jobs:
-                            # show error name and error message (if any)
-                            jobs.error(
-                                self.target.table_name, self._job_key(key),
-                                error_message=error_message, error_stack=traceback.format_exc())
-                        if not suppress_errors or isinstance(error, SystemExit):
-                            raise
-                        else:
-                            logger.error(error)
-                            error_list.append((key, error if return_exception_objects else error_message))
-                    else:
-                        self.connection.commit_transaction()
-                        if reserve_jobs:
-                            jobs.complete(self.target.table_name, self._job_key(key))
-                    finally:
-                        self.__class__._allow_insert = False
+                    for error in pool.imap(_call_populate1, keys):
+                        if error is not None:
+                            error_list.append(error)
+            self.connection.connect()  # reconnect parent process to MySQL server
 
-        # place back the original signal handler
+        # restore original signal handler:
         if reserve_jobs:
             signal.signal(signal.SIGTERM, old_handler)
-        return error_list
+
+        if suppress_errors:
+            return error_list
+
+    def _populate1(self, key, jobs, suppress_errors, return_exception_objects):
+        """
+        populates table for one source key, calling self.make inside a transaction.
+        :param jobs: the jobs table or None if not reserve_jobs
+        :param key: dict specifying job to populate
+        :param suppress_errors: bool if errors should be suppressed and returned
+        :param return_exception_objects: if True, errors must be returned as objects
+        :return: (key, error) when suppress_errors=True, otherwise None
+        """
+        make = self._make_tuples if hasattr(self, '_make_tuples') else self.make
+
+        if jobs is None or jobs.reserve(self.target.table_name, self._job_key(key)):
+            self.connection.start_transaction()
+            if key in self.target:  # already populated
+                self.connection.cancel_transaction()
+                if jobs is not None:
+                    jobs.complete(self.target.table_name, self._job_key(key))
+            else:
+                logger.info('Populating: ' + str(key))
+                self.__class__._allow_insert = True
+                try:
+                    make(dict(key))
+                except (KeyboardInterrupt, SystemExit, Exception) as error:
+                    try:
+                        self.connection.cancel_transaction()
+                    except LostConnectionError:
+                        pass
+                    error_message = '{exception}{msg}'.format(
+                        exception=error.__class__.__name__,
+                        msg=': ' + str(error) if str(error) else '')
+                    if jobs is not None:
+                        # show error name and error message (if any)
+                        jobs.error(
+                            self.target.table_name, self._job_key(key),
+                            error_message=error_message, error_stack=traceback.format_exc())
+                    if not suppress_errors or isinstance(error, SystemExit):
+                        raise
+                    else:
+                        logger.error(error)
+                        return key, error if return_exception_objects else error_message
+                else:
+                    self.connection.commit_transaction()
+                    if jobs is not None:
+                        jobs.complete(self.target.table_name, self._job_key(key))
+                finally:
+                    self.__class__._allow_insert = False
 
     def progress(self, *restrictions, display=True):
         """
-        report progress of populating the table
-        :return: remaining, total -- tuples to be populated
+        Report the progress of populating the table.
+        :return: (remaining, total) -- numbers of tuples to be populated
         """
         todo = self._jobs_to_do(restrictions)
         total = len(todo)
@@ -194,5 +270,6 @@ def progress(self, *restrictions, display=True):
             print('%-20s' % self.__class__.__name__,
                   'Completed %d of %d (%2.1f%%)   %s' % (
                       total - remaining, total, 100 - 100 * remaining / (total+1e-12),
-                      datetime.datetime.strftime(datetime.datetime.now(), '%Y-%m-%d %H:%M:%S')), flush=True)
+                      datetime.datetime.strftime(datetime.datetime.now(),
+                                                 '%Y-%m-%d %H:%M:%S')), flush=True)
         return remaining, total

datajoint/blob.py

@@ -166,6 +166,8 @@ def pack_blob(self, obj):
             return self.pack_array(np.array(obj))
         if isinstance(obj, (bool, np.bool_)):
             return self.pack_array(np.array(obj))
+        if isinstance(obj, (float, int, complex)):
+            return self.pack_array(np.array(obj))
         if isinstance(obj, (datetime.datetime, datetime.date, datetime.time)):
             return self.pack_datetime(obj)
         if isinstance(obj, Decimal):

datajoint/connection.py

@@ -278,7 +278,7 @@ def query(self, query, args=(), *, as_dict=False, suppress_warnings=True, reconn
         # check cache first:
         use_query_cache = bool(self._query_cache)
         if use_query_cache and not re.match(r"\s*(SELECT|SHOW)", query):
-            raise errors.DataJointError("Only SELECT query are allowed when query caching is on.")
+            raise errors.DataJointError("Only SELECT queries are allowed when query caching is on.")
         if use_query_cache:
             if not config['query_cache']:
                 raise errors.DataJointError("Provide filepath dj.config['query_cache'] when using query caching.")

datajoint/dependencies.py

@@ -26,8 +26,6 @@ def unite_master_parts(lst):
                     # move from the ith position to the (j+1)th position
                     lst[j+1:i+1] = [name] + lst[j+1:i]
                     break
-            else:
-                raise DataJointError("Found a part table {name} without its master table.".format(name=name))
     return lst