#11878

pending completion

Build # #11878

Build Type

push

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/datajoint/datajoint-python/commit/<a class=hub.com/datajoint/datajoint-python/commit/83e7922e5f74f739e87f6cba70c883a16bac8060">83e7922e5<a href="https://github.com/datajoint/datajoint-python/commit/83e7922e5f74f739e87f6cba70c883a16bac8060">">Merge </a><a class="double-link" href="https://github.com/datajoint/datajoint-python/commit/<a class="double-link" href="https://github.com/datajoint/datajoint-python/commit/f86de871d320815f7800812aa8ce4158ade2ef23">f86de871d</a>">f86de871d</a><a href="https://github.com/datajoint/datajoint-python/commit/83e7922e5f74f739e87f6cba70c883a16bac8060"> into e67046e4c">e67046e4c</a>

Run Details

1 of 1 new or added line in 1 file covered. (100.0%)

3053 of 3386 relevant lines covered (90.17%)

0.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

72.79

/datajoint/autopopulate.py

"""This module defines class dj.AutoPopulate"""
import logging
import datetime
import traceback
import random
import inspect
from tqdm import tqdm
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.
    Auto-populated relations must inherit from both Relation and AutoPopulate,
    must define the property `key_source`, and must define the callback method `make`.
    """
    _key_source = None
    _allow_insert = False

    @property
    def key_source(self):
        """
        :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.proj())

        if self._key_source is None:
            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')
            self._key_source = _rename_attributes(*parents[0])
            for q in parents[1:]:
                self._key_source *= _rename_attributes(*q)
        return self._key_source

    def make(self, key):
        """
        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`')

    @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.
        """
        return self

    def _job_key(self, key):
        """
        :param key:  they key returned for the job from the key source
        :return: the dict to use to generate the job reservation hash
        This method allows subclasses to control the job reservation granularity.
        """
        return key

    def _jobs_to_do(self, restrictions):
        """
        :return: the relation containing the keys to be computed (derived from self.key_source)
        """
        if self.restriction:
            raise DataJointError('Cannot call populate on a restricted table. '
                                 'Instead, pass conditions to populate() as arguments.')
        todo = self.key_source

        # key_source is a QueryExpression subclass -- trigger instantiation
        if inspect.isclass(todo) and issubclass(todo, QueryExpression):
            todo = todo()

        if not isinstance(todo, QueryExpression):
            raise DataJointError('Invalid key_source value')

        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))
        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, processes=1, make_kwargs=None):
        """
        ``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, 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 processes: number of processes to use. When set to a large number, then
            uses as many as CPU cores
        :param make_kwargs: Keyword arguments which do not affect the result of computation
            to be passed down to each ``make()`` call. Computation arguments should be
            specified within the pipeline e.g. using a `dj.Lookup` table.
        :type make_kwargs: dict, optional
        """
        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))
        jobs = self.connection.schemas[self.target.database].jobs if reserve_jobs else None

        # define and set up signal handler for SIGTERM:
        if reserve_jobs:
            def handler(signum, frame):
                logger.info('Populate terminated by SIGTERM')
                raise SystemExit('SIGTERM received')
            old_handler = signal.signal(signal.SIGTERM, handler)

        keys = (self._jobs_to_do(restrictions) - self.target).fetch("KEY", limit=limit)
        if order == "reverse":
            keys.reverse()
        elif order == "random":
            random.shuffle(keys)

        logger.info('Found %d keys to populate' % len(keys))

        keys = keys[:max_calls]
        nkeys = len(keys)

        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,
            make_kwargs=make_kwargs)

        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, jobs, 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:
                    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

        # restore original signal handler:
        if reserve_jobs:
            signal.signal(signal.SIGTERM, old_handler)

        if suppress_errors:
            return error_list

    def _populate1(self, key, jobs, suppress_errors, return_exception_objects, make_kwargs=None):
        """
        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), **(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 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 the progress of populating the table.
        :return: (remaining, total) -- numbers of tuples to be populated
        """
        todo = self._jobs_to_do(restrictions)
        total = len(todo)
        remaining = len(todo - self.target)
        if display:
            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)
        return remaining, total

1	"""This module defines class dj.AutoPopulate"""
2	import logging	1✔
3	import datetime	1✔
4	import traceback	1✔
5	import random	1✔
6	import inspect	1✔
7	from tqdm import tqdm	1✔
8	from .expression import QueryExpression, AndList	1✔
9	from .errors import DataJointError, LostConnectionError	1✔
10	import signal	1✔
11	import multiprocessing as mp	1✔
12
13	# noinspection PyExceptionInherit,PyCallingNonCallable
14
15	logger = logging.getLogger(__name__)	1✔
16
17
18	# --- helper functions for multiprocessing --
19
20	def _initialize_populate(table, jobs, populate_kwargs):	1✔
21	"""
22	Initialize the process for mulitprocessing.
23	Saves the unpickled copy of the table to the current process and reconnects.
24	"""
25	process = mp.current_process()	×
26	process.table = table	×
27	process.jobs = jobs	×
28	process.populate_kwargs = populate_kwargs	×
29	table.connection.connect() # reconnect	×
30
31
32	def _call_populate1(key):	1✔
33	"""
34	Call current process' table._populate1()
35	:key - a dict specifying job to compute
36	:return: key, error if error, otherwise None
37	"""
38	process = mp.current_process()	×
39	return process.table._populate1(key, process.jobs, **process.populate_kwargs)	×
40
41
42	class AutoPopulate:	1✔
43	"""
44	AutoPopulate is a mixin class that adds the method populate() to a Relation class.
45	Auto-populated relations must inherit from both Relation and AutoPopulate,
46	must define the property `key_source`, and must define the callback method `make`.
47	"""
48	_key_source = None	1✔
49	_allow_insert = False	1✔
50
51	@property	1✔
52	def key_source(self):	1✔
53	"""
54	:return: the query expression that yields primary key values to be passed,
55	sequentially, to the ``make`` method when populate() is called.
56	The default value is the join of the parent tables references from the primary key.
57	Subclasses may override they key_source to change the scope or the granularity
58	of the make calls.
59	"""
60	def _rename_attributes(table, props):	1✔
61	return (table.proj(	1✔
62	**{attr: ref for attr, ref in props['attr_map'].items() if attr != ref})
63	if props['aliased'] else table.proj())
64
65	if self._key_source is None:	1✔
66	parents = self.target.parents(	1✔
67	primary=True, as_objects=True, foreign_key_info=True)
68	if not parents:	1✔
69	raise DataJointError('A table must have dependencies '	×
70	'from its primary key for auto-populate to work')
71	self._key_source = _rename_attributes(*parents[0])	1✔
72	for q in parents[1:]:	1✔
73	self._key_source = _rename_attributes(q)	1✔
74	return self._key_source	1✔
75
76	def make(self, key):	1✔
77	"""
78	Derived classes must implement method `make` that fetches data from tables
79	above them in the dependency hierarchy, restricting by the given key,
80	computes secondary attributes, and inserts the new tuples into self.
81	"""
82	raise NotImplementedError(	×
83	'Subclasses of AutoPopulate must implement the method `make`')
84
85	@property	1✔
86	def target(self):	1✔
87	"""
88	:return: table to be populated.
89	In the typical case, dj.AutoPopulate is mixed into a dj.Table class by
90	inheritance and the target is self.
91	"""
92	return self	1✔
93
94	def _job_key(self, key):	1✔
95	"""
96	:param key: they key returned for the job from the key source
97	:return: the dict to use to generate the job reservation hash
98	This method allows subclasses to control the job reservation granularity.
99	"""
100	return key	1✔
101
102	def _jobs_to_do(self, restrictions):	1✔
103	"""
104	:return: the relation containing the keys to be computed (derived from self.key_source)
105	"""
106	if self.restriction:	1✔
107	raise DataJointError('Cannot call populate on a restricted table. '	×
108	'Instead, pass conditions to populate() as arguments.')
109	todo = self.key_source	1✔
110
111	# key_source is a QueryExpression subclass -- trigger instantiation
112	if inspect.isclass(todo) and issubclass(todo, QueryExpression):	1✔
113	todo = todo()	1✔
114
115	if not isinstance(todo, QueryExpression):	1✔
116	raise DataJointError('Invalid key_source value')	×
117
118	try:	1✔
119	# check if target lacks any attributes from the primary key of key_source
120	raise DataJointError(	1✔
121	'The populate target lacks attribute %s '
122	'from the primary key of key_source' % next(
123	name for name in todo.heading.primary_key
124	if name not in self.target.heading))
125	except StopIteration:	1✔
126	pass	1✔
127	return (todo & AndList(restrictions)).proj()	1✔
128
129	def populate(self, *restrictions, suppress_errors=False, return_exception_objects=False,	1✔
130	reserve_jobs=False, order="original", limit=None, max_calls=None,
131	display_progress=False, processes=1, make_kwargs=None):
132	"""
133	``table.populate()`` calls ``table.make(key)`` for every primary key in
134	``self.key_source`` for which there is not already a tuple in table.
135
136	:param restrictions: a list of restrictions each restrict
137	(table.key_source - target.proj())
138	:param suppress_errors: if True, do not terminate execution.
139	:param return_exception_objects: return error objects instead of just error messages
140	:param reserve_jobs: if True, reserve jobs to populate in asynchronous fashion
141	:param order: "original"\|"reverse"\|"random" - the order of execution
142	:param limit: if not None, check at most this many keys
143	:param max_calls: if not None, populate at most this many keys
144	:param display_progress: if True, report progress_bar
145	:param processes: number of processes to use. When set to a large number, then
146	uses as many as CPU cores
147	:param make_kwargs: Keyword arguments which do not affect the result of computation
148	to be passed down to each ``make()`` call. Computation arguments should be
149	specified within the pipeline e.g. using a `dj.Lookup` table.
150	:type make_kwargs: dict, optional
151	"""
152	if self.connection.in_transaction:	1✔
153	raise DataJointError('Populate cannot be called during a transaction.')	×
154
155	valid_order = ['original', 'reverse', 'random']	1✔
156	if order not in valid_order:	1✔
157	raise DataJointError('The order argument must be one of %s' % str(valid_order))	×
158	jobs = self.connection.schemas[self.target.database].jobs if reserve_jobs else None	1✔
159
160	# define and set up signal handler for SIGTERM:
161	if reserve_jobs:	1✔
162	def handler(signum, frame):	1✔
163	logger.info('Populate terminated by SIGTERM')	×
164	raise SystemExit('SIGTERM received')	×
165	old_handler = signal.signal(signal.SIGTERM, handler)	1✔
166
167	keys = (self._jobs_to_do(restrictions) - self.target).fetch("KEY", limit=limit)	1✔
168	if order == "reverse":	1✔
169	keys.reverse()	×
170	elif order == "random":	1✔
171	random.shuffle(keys)	×
172
173	logger.info('Found %d keys to populate' % len(keys))	1✔
174
175	keys = keys[:max_calls]	1✔
176	nkeys = len(keys)	1✔
177
178	if processes > 1:	1✔
179	processes = min(processes, nkeys, mp.cpu_count())	×
180
181	error_list = []	1✔
182	populate_kwargs = dict(	1✔
183	suppress_errors=suppress_errors,
184	return_exception_objects=return_exception_objects,
185	make_kwargs=make_kwargs)
186
187	if processes == 1:	1✔
188	for key in tqdm(keys, desc=self.__class__.__name__) if display_progress else keys:	1✔
189	error = self._populate1(key, jobs, **populate_kwargs)	1✔
190	if error is not None:	1✔
191	error_list.append(error)	1✔
192	else:
193	# spawn multiple processes
194	self.connection.close() # disconnect parent process from MySQL server	×
195	del self.connection._conn.ctx # SSLContext is not pickleable	×
196	with mp.Pool(processes, _initialize_populate, (self, jobs, populate_kwargs)) as pool:	×
197	if display_progress:	×
198	with tqdm(desc="Processes: ", total=nkeys) as pbar:	×
199	for error in pool.imap(_call_populate1, keys, chunksize=1):	×
200	if error is not None:	×
201	error_list.append(error)	×
202	pbar.update()	×
203	else:
204	for error in pool.imap(_call_populate1, keys):	×
205	if error is not None:	×
206	error_list.append(error)	×
207	self.connection.connect() # reconnect parent process to MySQL server	×
208
209	# restore original signal handler:
210	if reserve_jobs:	1✔
211	signal.signal(signal.SIGTERM, old_handler)	1✔
212
213	if suppress_errors:	1✔
214	return error_list	1✔
215
216	def _populate1(self, key, jobs, suppress_errors, return_exception_objects, make_kwargs=None):	1✔
217	"""
218	populates table for one source key, calling self.make inside a transaction.
219	:param jobs: the jobs table or None if not reserve_jobs
220	:param key: dict specifying job to populate
221	:param suppress_errors: bool if errors should be suppressed and returned
222	:param return_exception_objects: if True, errors must be returned as objects
223	:return: (key, error) when suppress_errors=True, otherwise None
224	"""
225	make = self._make_tuples if hasattr(self, '_make_tuples') else self.make	1✔
226
227	if jobs is None or jobs.reserve(self.target.table_name, self._job_key(key)):	1✔
228	self.connection.start_transaction()	1✔
229	if key in self.target: # already populated	1✔
230	self.connection.cancel_transaction()	×
231	if jobs is not None:	×
232	jobs.complete(self.target.table_name, self._job_key(key))	×
233	else:
234	logger.info('Populating: ' + str(key))	1✔
235	self.__class__._allow_insert = True	1✔
236	try:	1✔
237	make(dict(key), **(make_kwargs or {}))	1✔
238	except (KeyboardInterrupt, SystemExit, Exception) as error:	1✔
239	try:	1✔
240	self.connection.cancel_transaction()	1✔
241	except LostConnectionError:	×
242	pass	×
243	error_message = '{exception}{msg}'.format(	1✔
244	exception=error.__class__.__name__,
245	msg=': ' + str(error) if str(error) else '')
246	if jobs is not None:	1✔
247	# show error name and error message (if any)
248	jobs.error(	1✔
249	self.target.table_name, self._job_key(key),
250	error_message=error_message, error_stack=traceback.format_exc())
251	if not suppress_errors or isinstance(error, SystemExit):	1✔
252	raise	1✔
253	else:
254	logger.error(error)	1✔
255	return key, error if return_exception_objects else error_message	1✔
256	else:
257	self.connection.commit_transaction()	1✔
258	if jobs is not None:	1✔
259	jobs.complete(self.target.table_name, self._job_key(key))	×
260	finally:
261	self.__class__._allow_insert = False	1✔
262
263	def progress(self, *restrictions, display=True):	1✔
264	"""
265	Report the progress of populating the table.
266	:return: (remaining, total) -- numbers of tuples to be populated
267	"""
268	todo = self._jobs_to_do(restrictions)	1✔
269	total = len(todo)	1✔
270	remaining = len(todo - self.target)	1✔
271	if display:	1✔
272	print('%-20s' % self.__class__.__name__,	1✔
273	'Completed %d of %d (%2.1f%%) %s' % (
274	total - remaining, total, 100 - 100 * remaining / (total+1e-12),
275	datetime.datetime.strftime(datetime.datetime.now(),
276	'%Y-%m-%d %H:%M:%S')), flush=True)
277	return remaining, total	1✔

datajoint / datajoint-python / #11878

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous