bugfix, add tests

Author: Thinh Nguyen

datajoint/autopopulate.py

@@ -41,7 +41,9 @@ def _call_populate1(key):
     :return: key, error if error, otherwise None
     """
     process = mp.current_process()
-    return process.table._populate1(key, process.jobs, **process.populate_kwargs)
+    return process.table._populate1(
+        key, process.reserve_jobs, **process.populate_kwargs
+    )
 
 
 class AutoPopulate:
@@ -208,13 +210,18 @@ def handler(signum, frame):
             old_handler = signal.signal(signal.SIGTERM, handler)
 
             if schedule_jobs:
-                self.schedule_jobs()
+                self.schedule_jobs(*restrictions)
 
             keys = (
                 self._Jobs
                 & {"table_name": self.target.table_name}
                 & 'status = "scheduled"'
-            ).fetch("KEY", limit=limit)
+            ).fetch("key", limit=limit)
+
+            if restrictions:
+                # hitting the `key_source` again to apply the restrictions
+                # this is expensive/suboptimal
+                keys = (self._jobs_to_do(restrictions) & keys).fetch("KEY", limit=limit)
         else:
             keys = (self._jobs_to_do(restrictions) - self.target).fetch(
                 "KEY", limit=limit
@@ -338,6 +345,9 @@ def _populate1(
                             self._job_key(key),
                             error_message=error_message,
                             error_stack=traceback.format_exc(),
+                            run_duration=(
+                                datetime.datetime.utcnow() - make_start
+                            ).total_seconds(),
                         )
                     if not suppress_errors or isinstance(error, SystemExit):
                         raise
@@ -391,16 +401,18 @@ def _Jobs(self):
     def jobs(self):
         return self._Jobs & {"table_name": self.target.table_name}
 
-    def schedule_jobs(self, purge_invalid_jobs=True):
+    def schedule_jobs(self, *restrictions, purge_invalid_jobs=True):
         """
         Schedule new jobs for this autopopulate table
+        :param restrictions: a list of restrictions each restrict
+            (table.key_source - target.proj())
         :param purge_invalid_jobs: if True, remove invalid entry from the jobs table (potentially expensive operation)
         :return:
         """
         try:
             with self.connection.transaction:
                 schedule_count = 0
-                for key in (self._jobs_to_do({}) - self.target).fetch("KEY"):
+                for key in (self._jobs_to_do(restrictions) - self.target).fetch("KEY"):
                     schedule_count += self._Jobs.schedule(self.target.table_name, key)
         except Exception as e:
             logger.exception(str(e))
@@ -425,11 +437,11 @@ def purge_invalid_jobs(self):
         invalid_count = len(jobs_query) - len(self._jobs_to_do({}))
         invalid_removed = 0
         if invalid_count > 0:
-            for key, job_key in jobs_query.fetch("KEY", "key"):
+            for key, job_key in zip(*jobs_query.fetch("KEY", "key")):
                 if not (self._jobs_to_do({}) & job_key):
                     (jobs_query & key).delete()
                     invalid_removed += 1
 
-        logger.info(
-            f"{invalid_removed}/{invalid_count} invalid jobs removed for `{to_camel_case(self.target.table_name)}`"
-        )
+            logger.info(
+                f"{invalid_removed}/{invalid_count} invalid jobs removed for `{to_camel_case(self.target.table_name)}`"
+            )

datajoint/jobs.py

@@ -79,9 +79,9 @@ def schedule(self, table_name, key, seconds_delay=0, force=False):
         job_key = dict(table_name=table_name, key_hash=key_hash(key))
         if self & job_key:
             current_status = (self & job_key).fetch1("status")
-            if current_status in ("scheduled", "reserved", "success"):
-                return True
-            if current_status in ("error", "ignore") and not force:
+            if current_status in ("scheduled", "reserved", "success") or (
+                current_status in ("error", "ignore") and not force
+            ):
                 return False
 
         job = dict(
@@ -109,9 +109,14 @@ def reserve(self, table_name, key):
         :param key: the dict of the job's primary key
         :return: True if reserved job successfully. False = the jobs is already taken
         """
+        job_key = dict(table_name=table_name, key_hash=key_hash(key))
+        if self & job_key:
+            current_status = (self & job_key).fetch1("status")
+            if current_status != "scheduled":
+                return False
+
         job = dict(
-            table_name=table_name,
-            key_hash=key_hash(key),
+            job_key,
             status="reserved",
             host=platform.node(),
             pid=os.getpid(),
@@ -120,11 +125,10 @@ def reserve(self, table_name, key):
             user=self._user,
             timestamp=datetime.datetime.utcnow(),
         )
-        try:
-            with config(enable_python_native_blobs=True):
-                self.insert1(job, ignore_extra_fields=True)
-        except DuplicateError:
-            return False
+
+        with config(enable_python_native_blobs=True):
+            self.insert1(job, replace=True, ignore_extra_fields=True)
+
         return True
 
     def ignore(self, table_name, key, message=""):

datajoint/schemas.py

@@ -8,7 +8,7 @@
 from .diagram import Diagram, _get_tier
 from .settings import config
 from .errors import DataJointError, AccessError
-from .jobs import JobTable, JobConfigTable
+from .jobs import JobTable
 from .external import ExternalMapping
 from .heading import Heading
 from .utils import user_choice, to_camel_case

tests/test_autopopulate.py

@@ -24,6 +24,7 @@ def tearDown(self):
         self.trial.Condition.delete_quick()
         self.trial.delete_quick()
         self.experiment.delete_quick()
+        schema.schema.jobs.delete_quick()
 
     def test_populate(self):
         # test simple populate
@@ -70,6 +71,73 @@ def test_populate_with_success_count(self):
         )
         assert_equal(len(self.trial.key_source & self.trial), success_count)
 
+    def test_schedule_jobs(self):
+        assert_true(self.subject, "root tables are empty")
+        assert_false(self.experiment, "table already filled?")
+        # test schedule jobs
+        self.experiment.schedule_jobs()
+        assert_true(
+            len(
+                schema.schema.jobs
+                & {"table_name": self.experiment.table_name, "status": "scheduled"}
+            )
+            == len(self.experiment.key_source),
+            "failed scheduling jobs",
+        )
+        # test executing jobs
+        self.experiment.populate(reserve_jobs=True, schedule_jobs=False)
+        assert_true(
+            len(
+                schema.schema.jobs
+                & {"table_name": self.experiment.table_name, "status": "success"}
+            )
+            == len(self.experiment.key_source),
+            "failed executing jobs",
+        )
+        # test schedule and execute jobs with restriction
+        restriction = self.subject.proj(animal="subject_id").fetch("KEY")[0]
+        self.trial.schedule_jobs(restriction)
+        assert_true(
+            len(
+                schema.schema.jobs
+                & {"table_name": self.trial.table_name, "status": "scheduled"}
+            )
+            == len(self.trial.key_source & restriction),
+            "failed scheduling jobs",
+        )
+        self.trial.schedule_jobs()
+        assert_true(
+            len(
+                schema.schema.jobs
+                & {"table_name": self.trial.table_name, "status": "scheduled"}
+            )
+            == len(self.trial.key_source),
+            "failed scheduling jobs",
+        )
+        self.trial.populate(restriction, reserve_jobs=True, schedule_jobs=False)
+        assert_equal(
+            len(self.trial.key_source & self.trial),
+            len(self.trial.key_source & restriction),
+        )
+        assert_equal(
+            len(self.trial.key_source - self.trial),
+            len(self.trial.key_source - restriction),
+        )
+        self.trial.populate(reserve_jobs=True, schedule_jobs=False)
+        assert_equal(
+            len(self.trial.key_source & self.trial), len(self.trial.key_source)
+        )
+        # test purge invalid jobs
+        restriction["subject_id"] = restriction.pop("animal")
+        with dj.config(safemode=False):
+            (self.experiment & restriction).delete()
+        self.trial.purge_invalid_jobs()
+        assert_true(
+            len(schema.schema.jobs & {"table_name": self.trial.table_name})
+            == len(self.trial.key_source),
+            "failed purging invalid jobs",
+        )
+
     def test_populate_exclude_error_and_ignore_jobs(self):
         # test simple populate
         assert_true(self.subject, "root tables are empty")