Merge pull request datajoint#704 from mspacek/mp

Add multiprocessing to AutoPopulate

Author: dimitri-yatsenko

CHANGELOG.md

@@ -1,5 +1,10 @@
 ## Release notes
 
+## 0.12.3 -- Nov 22, 2019
+* Bugfix #675 (PR #705) networkx 2.4+ is now supported
+* Bugfix #698 and #699 (PR #706) display table definition in doc string and help
+* Bugfix #701 (PR #702) job reservation works with native python datatype support disabled
+
 ### 0.12.2 -- Nov 11, 2019
 * Bugfix - Convoluted error thrown if there is a reference to a non-existent table attribute (#691)
 * Bugfix - Insert into external does not trim leading slash if defined in `dj.config['stores']['<store>']['location']` (#692)

datajoint/autopopulate.py

@@ -10,12 +10,26 @@
 from .errors import DataJointError
 from .table import FreeTable
 import signal
+import multiprocessing as mp
 
 # noinspection PyExceptionInherit,PyCallingNonCallable
 
 logger = logging.getLogger(__name__)
 
 
+def initializer(table):
+    """Save pickled copy of (disconnected) table to the current process,
+    then reconnect to server. For use by call_make_key()"""
+    mp.current_process().table = table
+    table.connection.connect() # reconnect
+
+def call_make_key(key):
+    """Call current process' table.make_key()"""
+    table = mp.current_process().table
+    error = table.make_key(key)
+    return error
+
+
 class AutoPopulate:
     """
     AutoPopulate is a mixin class that adds the method populate() to a Relation class.
@@ -103,29 +117,36 @@ def _jobs_to_do(self, restrictions):
 
     def populate(self, *restrictions, suppress_errors=False, return_exception_objects=False,
                  reserve_jobs=False, order="original", limit=None, max_calls=None,
-                 display_progress=False):
+                 display_progress=False, multiprocess=False):
         """
         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())
         :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 multiprocess: if True, use as many processes as CPU cores, or use the integer
+        number of processes specified
         """
         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
+        self._make_key_kwargs = {'suppress_errors':suppress_errors,
+                                 'return_exception_objects':return_exception_objects,
+                                 'reserve_jobs':reserve_jobs,
+                                 'jobs':jobs,
+                                }
+
+        # define and set up signal handler for SIGTERM:
         if reserve_jobs:
             def handler(signum, frame):
                 logger.info('Populate terminated by SIGTERM')
@@ -138,55 +159,101 @@ 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
+        if max_calls is not None:
+            keys = keys[:max_calls]
+        nkeys = len(keys)
 
-        for key in (tqdm(keys) 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 multiprocess: # True or int, presumably
+            if multiprocess is True:
+                nproc = mp.cpu_count()
+            else:
+                if not isinstance(multiprocess, int):
+                    raise DataJointError("multiprocess can be False, True or a positive integer")
+                nproc = multiprocess
+        else:
+            nproc = 1
+        nproc = min(nproc, nkeys) # no sense spawning more than can be used
+        error_list = []
+        if nproc > 1: # spawn multiple processes
+            # prepare to pickle self:
+            self.connection.close() # disconnect parent process from MySQL server
+            del self.connection._conn.ctx # SSLContext is not picklable
+            print('*** Spawning pool of %d processes' % nproc)
+            # send pickled copy of self to each process,
+            # each worker process calls initializer(*initargs) when it starts
+            with mp.Pool(nproc, initializer, (self,)) as pool:
+                if display_progress:
+                    with tqdm(total=nkeys) as pbar:
+                        for error in pool.imap(call_make_key, 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))
-                    except (KeyboardInterrupt, SystemExit, Exception) as error:
-                        try:
-                            self.connection.cancel_transaction()
-                        except OperationalError:
-                            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_make_key, keys):
+                        if error is not None:
+                            error_list.append(error)
+            self.connection.connect() # reconnect parent process to MySQL server
+        else: # use single process
+            for key in tqdm(keys) if display_progress else keys:
+                error = self.make_key(key)
+                if error is not None:
+                    error_list.append(error)
 
-        # place back the original signal handler
+        del self._make_key_kwargs # clean up
+
+        # restore original signal handler:
         if reserve_jobs:
             signal.signal(signal.SIGTERM, old_handler)
-        return error_list
+
+        if suppress_errors:
+            return error_list
+
+    def make_key(self, key):
+        make = self._make_tuples if hasattr(self, '_make_tuples') else self.make
+
+        kwargs = self._make_key_kwargs
+        suppress_errors = kwargs['suppress_errors']
+        return_exception_objects = kwargs['return_exception_objects']
+        reserve_jobs = kwargs['reserve_jobs']
+        jobs = kwargs['jobs']
+
+        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))
+            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 OperationalError:
+                        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)
+                        return (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
 
     def progress(self, *restrictions, display=True):
         """

datajoint/blob.py

@@ -152,10 +152,8 @@ def pack_blob(self, obj):
             return self.pack_array(np.array(obj))
         if isinstance(obj, (bool, np.bool, np.bool_)):
             return self.pack_array(np.array(obj))
-        if isinstance(obj, float):
-            return self.pack_array(np.array(obj, dtype=np.float64))
-        if isinstance(obj, int):
-            return self.pack_array(np.array(obj, dtype=np.int64))
+        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/heading.py

@@ -247,7 +247,10 @@ def init_from_database(self, conn, database, table_name, context):
                     category = next(c for c in SPECIAL_TYPES if TYPE_PATTERN[c].match(attr['type']))
                 except StopIteration:
                     if attr['type'].startswith('external'):
-                        raise DataJointError('Legacy datatype `{type}`.'.format(**attr)) from None
+                        url = "https://docs.datajoint.io/python/admin/5-blob-config.html" \
+                              "#migration-between-datajoint-v0-11-and-v0-12"
+                        raise DataJointError('Legacy datatype `{type}`. Migrate your external stores to '
+                                             'datajoint 0.12: {url}'.format(url=url, **attr)) from None
                     raise DataJointError('Unknown attribute type `{type}`'.format(**attr)) from None
                 if category == 'FILEPATH' and not _support_filepath_types():
                     raise DataJointError("""

datajoint/schema.py

@@ -191,6 +191,10 @@ def process_table_class(self, table_class, context, assert_declared=False):
                 instance.declare(context)
         is_declared = is_declared or instance.is_declared
 
+        # add table definition to the doc string
+        if isinstance(table_class.definition, str):
+            table_class.__doc__ = (table_class.__doc__ or "") + "\nTable definition:\n\n" + table_class.definition
+
         # fill values in Lookup tables from their contents property
         if isinstance(instance, Lookup) and hasattr(instance, 'contents') and is_declared:
             contents = list(instance.contents)

datajoint/table.py

@@ -45,14 +45,9 @@ def heading(self):
         """
         if self._heading is None:
             self._heading = Heading()  # instance-level heading
-        if not self._heading:  # lazy loading of heading
-            if self.connection is None:
-                raise DataJointError(
-                    'DataJoint class is missing a database connection. '
-                    'Missing schema decorator on the class? (e.g. @schema)')
-            else:
-                self._heading.init_from_database(
-                    self.connection, self.database, self.table_name, self.declaration_context)
+        if not self._heading and self.connection is not None:  # lazy loading of heading
+            self._heading.init_from_database(
+                self.connection, self.database, self.table_name, self.declaration_context)
         return self._heading
 
     def declare(self, context=None):
@@ -411,7 +406,7 @@ def delete(self, verbose=True):
             print('About to delete:')
 
         if not already_in_transaction:
-            self.connection.start_transaction()
+            conn.start_transaction()
         total = 0
         try:
             for name, table in reversed(list(delete_list.items())):
@@ -423,25 +418,25 @@ def delete(self, verbose=True):
         except:
             # Delete failed, perhaps due to insufficient privileges. Cancel transaction.
             if not already_in_transaction:
-                self.connection.cancel_transaction()
+                conn.cancel_transaction()
             raise
         else:
             assert not (already_in_transaction and safe)
             if not total:
                 print('Nothing to delete')
                 if not already_in_transaction:
-                    self.connection.cancel_transaction()
+                    conn.cancel_transaction()
             else:
                 if already_in_transaction:
                     if verbose:
                         print('The delete is pending within the ongoing transaction.')
                 else:
                     if not safe or user_choice("Proceed?", default='no') == 'yes':
-                        self.connection.commit_transaction()
+                        conn.commit_transaction()
                         if verbose or safe:
                             print('Committed.')
                     else:
-                        self.connection.cancel_transaction()
+                        conn.cancel_transaction()
                         if verbose or safe:
                             print('Cancelled deletes.')
 

datajoint/user_tables.py

@@ -13,7 +13,7 @@
 # attributes that trigger instantiation of user classes
 supported_class_attrs = {
     'key_source', 'describe', 'alter', 'heading', 'populate', 'progress', 'primary_key', 'proj', 'aggr',
-    'fetch', 'fetch1','head', 'tail',
+    'fetch', 'fetch1', 'head', 'tail',
     'insert', 'insert1', 'drop', 'drop_quick', 'delete', 'delete_quick'}
 
 
@@ -92,7 +92,7 @@ def table_name(cls):
 
     @ClassProperty
     def full_table_name(cls):
-        if cls not in {Manual, Imported, Lookup, Computed, Part}:
+        if cls not in {Manual, Imported, Lookup, Computed, Part, UserTable}:
             if cls.database is None:
                 raise DataJointError('Class %s is not properly declared (schema decorator not applied?)' % cls.__name__)
             return r"`{0:s}`.`{1:s}`".format(cls.database, cls.table_name)

datajoint/version.py

@@ -1,3 +1,3 @@
-__version__ = "0.12.2"
+__version__ = "0.12.3"
 
 assert len(__version__) <= 10  # The log table limits version to the 10 characters

docs-parts/intro/Releases_lang1.rst

@@ -1,4 +1,11 @@
-0.12.1 -- Nov 11, 2019
+0.12.3 -- Nov 22, 2019
+----------------------
+* Bugfix - networkx 2.4 causes error in diagrams (#675) PR #705
+* Bugfix - include definition in doc string and help (#698, #699) PR #706
+* Bugfix - job reservation fails when native python datatype support is disabled (#701) PR #702 
+
+
+0.12.2 -- Nov 11, 2019
 -------------------------
 * Bugfix - Convoluted error thrown if there is a reference to a non-existent table attribute (#691)
 * Bugfix - Insert into external does not trim leading slash if defined in `dj.config['stores']['<store>']['location']` (#692)

tests/schema.py

@@ -13,6 +13,9 @@
 
 @schema
 class TTest(dj.Lookup):
+    """
+    doc string
+    """
     definition = """
     key   :   int     # key
     ---