#12173

pending completion

Build # #12173

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/b02363992f7dcaf0ddc2ec1ef006b0934feb5008">b02363992<a href="https://github.com/datajoint/datajoint-python/commit/b02363992f7dcaf0ddc2ec1ef006b0934feb5008">">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/e8652f799626b321b7e7029804f9f83d62c845d4">e8652f799</a>">e8652f799</a><a href="https://github.com/datajoint/datajoint-python/commit/b02363992f7dcaf0ddc2ec1ef006b0934feb5008"> into e2aa1ad38">e2aa1ad38</a>

Run Details

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

2989 of 3316 relevant lines covered (90.14%)

0.9 hits per line

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

77.37

/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 < 0:
            raise Exception("processes must not be negative")
        elif processes == 0:
            return error_list, nkeys
        elif 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)

        return error_list, nkeys

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

datajoint / datajoint-python / #12173

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