#12817

pending completion

Build # #12817

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/6cae2c645e610cdc959c5e4f627a09d6e4fa00e0">6cae2c645<a href="https://github.com/datajoint/datajoint-python/commit/6cae2c645e610cdc959c5e4f627a09d6e4fa00e0">">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/37801d69847287e148a17191a423ea10e52e02a1">37801d698</a>">37801d698</a><a href="https://github.com/datajoint/datajoint-python/commit/6cae2c645e610cdc959c5e4f627a09d6e4fa00e0"> into a745322e1">a745322e1</a>

Run Details

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

3104 of 3427 relevant lines covered (90.57%)

0.91 hits per line

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

82.99

/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
import contextlib

# noinspection PyExceptionInherit,PyCallingNonCallable

logger = logging.getLogger(__name__.split(".")[0])


# --- 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 Table class.
    Auto-populated tables must inherit from both Table 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 query yeilding 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,
        return_success_count=False,
        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. Set to None to use all cores
        :param return_success_count: if True, return the count of successful `make()` calls.
            If suppress_errors is also True, returns a tuple: (success_count, errors)
        :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.debug("Found %d keys to populate" % len(keys))

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

        error_list = []
        success_list = []

        if nkeys:
            processes = min(_ for _ in (processes, nkeys, mp.cpu_count()) if _)

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

datajoint / datajoint-python / #12817

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