#12861

pending completion

Build # #12861

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/868947ab8f415b307f376d0f9b397800c900b58a">868947ab8<a href="https://github.com/datajoint/datajoint-python/commit/868947ab8f415b307f376d0f9b397800c900b58a">">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/7a258d438391b2ca5f3d7a7942cf8ab460f5cf72">7a258d438</a>">7a258d438</a><a href="https://github.com/datajoint/datajoint-python/commit/868947ab8f415b307f376d0f9b397800c900b58a"> into e339d465f">e339d465f</a>

Run Details

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

3036 of 3358 relevant lines covered (90.41%)

0.9 hits per line

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

83.89

/datajoint/autopopulate.py

"""This module defines class dj.AutoPopulate"""
import logging
import datetime
import traceback
import random
import inspect
from tqdm import tqdm
from .hash import key_hash
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)

        # exclude "error" or "ignore" jobs
        if reserve_jobs:
            exclude_key_hashes = (
                jobs
                & {"table_name": self.target.table_name}
                & 'status in ("error", "ignore")'
            ).fetch("key_hash")
            keys = [key for key in keys if key_hash(key) not in exclude_key_hashes]

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

datajoint / datajoint-python / #12861

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