diff --git a/src/datajoint/jobs.py b/src/datajoint/jobs.py index f2bc8d39b..3b23497d4 100644 --- a/src/datajoint/jobs.py +++ b/src/datajoint/jobs.py @@ -20,6 +20,7 @@ ERROR_MESSAGE_LENGTH = 2047 TRUNCATION_APPENDIX = "...truncated" +VANISHED_DELETE_BATCH = 1000 # job keys compared and deleted per page when refresh() removes vanished jobs logger = logging.getLogger(__name__.split(".")[0]) @@ -423,24 +424,21 @@ def refresh( # 3. Remove jobs whose keys have vanished. Judge against the unrestricted key_source: # jobs outside the restrictions (e.g. reserved by other workers) are not stale. # Disable semantic_check for Job table operations (job table PK has different lineage than target) - vanished = self.restrict(Not(self._target.key_source.proj()), semantic_check=False) # 3a. Success and error records describe rows that no longer exist once their key # has left both key_source and the target, regardless of age. - result["removed"] += ( - (vanished & "status in ('success', 'error')") - .restrict(Not(self._target.proj()), semantic_check=False) - .delete_quick(get_count=True) + result["removed"] += self._delete_vanished( + (self & "status in ('success', 'error')").restrict(Not(self._target.proj()), semantic_check=False) ) # 3b. Remove stale jobs (not ignore status) once older than stale_timeout, a grace period # for a key_source that is still being filled. stale_timeout=0 means no grace period. # Age is judged on the server's CURRENT_TIMESTAMP for consistent timing. - stale_jobs = vanished & "status != 'ignore'" + stale_jobs = self & "status != 'ignore'" if stale_timeout > 0: stale_interval = self.adapter.interval_expr(stale_timeout, "second") stale_jobs = stale_jobs & f"created_time < CURRENT_TIMESTAMP - {stale_interval}" - result["removed"] += stale_jobs.delete_quick(get_count=True) + result["removed"] += self._delete_vanished(stale_jobs) # 4. Handle orphaned reserved jobs - use server CURRENT_TIMESTAMP for consistent timing if orphan_timeout is not None and orphan_timeout > 0: @@ -454,6 +452,63 @@ def refresh( return result + def _delete_vanished(self, candidates) -> int: + """ + Delete the jobs in ``candidates`` whose keys are no longer in key_source. + + ``key_source`` is queried at the top level, restricted to the candidates, and + never nested in ``NOT IN``: MySQL materializes a ``NOT IN`` subquery, and a + ``key_source`` restricted by several part tables is then built as the join of + those parts, whose size is the product of their rows. Most refreshes find + nothing vanished, which two counts establish without fetching keys. + + Otherwise the candidates are compared with ``key_source`` one page of + ``VANISHED_DELETE_BATCH`` keys at a time, so memory stays bounded however + many jobs the table holds. Each page's vanished jobs are deleted in one + statement, restricted again by ``candidates``, so a job that has left the + candidates since the page was read is kept. + + Parameters + ---------- + candidates : QueryExpression + A restriction of this jobs table. + + Returns + ------- + int + Number of jobs deleted. + """ + n_candidates = len(candidates) + if not n_candidates: + return 0 + key_source = self._target.key_source.proj() + same_key = set(self.primary_key) == set(key_source.primary_key) + if same_key and len(key_source.restrict(candidates.proj(), semantic_check=False)) == n_candidates: + return 0 + + shared = [name for name in self.primary_key if name in key_source.primary_key] + deleted = kept = 0 + while True: + # Pages are ordered by key and deletions only remove keys before the + # cursor, so the next page starts after the keys left in place so far. + page = candidates.keys(order_by="KEY", limit=VANISHED_DELETE_BATCH, offset=kept) + if not page: + break + in_page = candidates & page + present = { + tuple(key[name] for name in shared) for key in key_source.restrict(in_page.proj(), semantic_check=False).keys() + } + vanished = [key for key in page if tuple(key[name] for name in shared) not in present] + n_deleted = (candidates & vanished).delete_quick(get_count=True) if vanished else 0 + deleted += n_deleted + # Advance by the rows actually left in place. A key that does not match its + # own value as a restriction, such as an inexact float, is neither found + # present nor deleted; counting it as kept moves the cursor past it. + kept += len(page) - n_deleted + if len(page) < VANISHED_DELETE_BATCH: + break + return deleted + def reserve(self, key: dict) -> bool: """ Attempt to reserve a pending job for processing. diff --git a/tests/integration/test_jobs_refresh.py b/tests/integration/test_jobs_refresh.py index c339063e2..e0a485f8d 100644 --- a/tests/integration/test_jobs_refresh.py +++ b/tests/integration/test_jobs_refresh.py @@ -201,7 +201,7 @@ def test_configured_stale_timeout_zero_means_no_grace_period(pipeline): def test_refresh_cleanup_is_set_based(pipeline, monkeypatch): """Cleanup issues a fixed number of statements, not one per vanished key.""" Raw, Proc = pipeline - n = 20 + n = 50 # more vanished keys than refresh() issues statements in all with dj.config.override(jobs={"keep_completed": True}): Raw.insert([dict(raw_id=i, v=float(i)) for i in range(n)]) Proc.populate(reserve_jobs=True) diff --git a/tests/integration/test_jobs_refresh_part_restricted.py b/tests/integration/test_jobs_refresh_part_restricted.py new file mode 100644 index 000000000..6890feaee --- /dev/null +++ b/tests/integration/test_jobs_refresh_part_restricted.py @@ -0,0 +1,356 @@ +""" +Cleanup in ``Job.refresh()`` for a ``key_source`` restricted by part tables. + +A ``key_source`` such as ``A * B & A.Item & B.Item`` keeps the keys whose masters +carry parts. Checking a job against it with ``NOT IN (key_source)`` lets MySQL +materialize the subquery as a join of both part tables, whose size is the product +of their rows: with tens of thousands of parts per master, every ``refresh()`` +took minutes, and a loop that refreshes before every ``populate()`` paid that on +every pass. The cleanup checks the jobs against ``key_source`` without nesting it +in ``NOT IN``. +""" + +import time +from datetime import datetime + +import pytest + +import datajoint as dj + +N_A_ITEMS = 3000 +N_B_ITEMS = 2000 + + +@pytest.fixture(scope="function") +def schema_by_backend(connection_by_backend, db_creds_by_backend): + backend = db_creds_by_backend["backend"] + test_id = str(int(time.time() * 1000))[-8:] + schema_name = f"djtest_partks_{backend}_{test_id}"[:64] + drop = connection_by_backend.adapter.drop_schema_sql(schema_name) + connection_by_backend.query(drop) + schema = dj.Schema(schema_name, connection=connection_by_backend) + yield schema + if connection_by_backend.is_connected: + connection_by_backend.query(drop) + + +@pytest.fixture +def pipeline(schema_by_backend): + @schema_by_backend + class Session(dj.Manual): + definition = """ + session_id : int32 + """ + + @schema_by_backend + class Profile(dj.Manual): + definition = """ + -> Session + profile_id : int16 + """ + + class Item(dj.Part): + definition = """ + -> master + item : int32 + """ + + @schema_by_backend + class Grouping(dj.Manual): + definition = """ + -> Session + grouping_id : int16 + """ + + class Item(dj.Part): + definition = """ + -> master + item : int32 + """ + + @schema_by_backend + class Figure(dj.Computed): + definition = """ + -> Profile + -> Grouping + --- + n : int32 + """ + + @property + def key_source(self): + # both masters must carry parts, not just keys + return Profile * Grouping & Profile.Item & Grouping.Item + + def make(self, key): + self.insert1(dict(key, n=len(Profile.Item & key))) + + return Session, Profile, Grouping, Figure + + +def _add(master, key, n_items): + master.insert1(key) + master.Item.insert([dict(key, item=i) for i in range(n_items)]) + + +def _backdate(jobs): + """Make every job older than the default stale_timeout.""" + for key in jobs.keys(): + jobs.update1({**key, "created_time": datetime(2000, 1, 1)}) + + +def _statuses(jobs): + return {row["profile_id"]: row["status"] for row in jobs.proj("status").to_dicts()} + + +def test_refresh_removes_jobs_of_vanished_keys(pipeline): + """Records of a deleted key go at once; a key whose parts went goes after stale_timeout.""" + Session, Profile, Grouping, Figure = pipeline + with dj.config.override(jobs={"keep_completed": True}): + Session.insert1(dict(session_id=1)) + _add(Grouping, dict(session_id=1, grouping_id=0), 3) + for p in range(3): + _add(Profile, dict(session_id=1, profile_id=p), 3) + Figure.populate(reserve_jobs=True) + assert _statuses(Figure.jobs) == {0: "success", 1: "success", 2: "success"} + + (Profile & {"profile_id": 0}).delete(prompt=False) # the row and its key go + (Profile.Item & {"profile_id": 1}).delete_quick() # the key leaves key_source; the row stays + + result = Figure.jobs.refresh() + assert _statuses(Figure.jobs) == {1: "success", 2: "success"} + assert result["removed"] == 1 + + _backdate(Figure.jobs) + result = Figure.jobs.refresh() + assert _statuses(Figure.jobs) == {2: "success"} + assert result["removed"] == 1 + + +def test_refresh_does_not_join_the_parts(pipeline, connection_by_backend, db_creds_by_backend): + """refresh() stays fast when key_source is restricted by two large part tables.""" + if db_creds_by_backend["backend"] != "mysql": + pytest.skip("MySQL's semijoin materialization") + Session, Profile, Grouping, Figure = pipeline + conn = connection_by_backend + with dj.config.override(jobs={"keep_completed": True}): + Session.insert1(dict(session_id=1)) + _add(Grouping, dict(session_id=1, grouping_id=0), N_B_ITEMS) + for p in range(2): + _add(Profile, dict(session_id=1, profile_id=p), N_A_ITEMS) + Figure.populate(reserve_jobs=True) + _backdate(Figure.jobs) + + # At these sizes MySQL may still pick FirstMatch; with tens of thousands of + # parts it picks materialization, which is forced here to keep the test small. + switch = conn.query("SELECT @@SESSION.optimizer_switch").fetchone()[0] + conn.query("SET SESSION optimizer_switch='firstmatch=off,loosescan=off,duplicateweedout=off'") + try: + start = time.monotonic() + result = Figure.jobs.refresh() + elapsed = time.monotonic() - start + finally: + conn.query(f"SET SESSION optimizer_switch='{switch}'") + + assert result["removed"] == 0 + assert len(Figure.jobs.completed) == 2 + assert elapsed < 2, f"refresh() took {elapsed:.1f} s" + + +@pytest.fixture +def one_master(schema_by_backend): + """A master with two part tables; key_source requires rows in both.""" + + @schema_by_backend + class Master(dj.Manual): + definition = """ + master_id : int32 + """ + + class PartA(dj.Part): + definition = """ + -> master + a : int32 + """ + + class PartB(dj.Part): + definition = """ + -> master + b : int32 + """ + + @schema_by_backend + class Summary(dj.Computed): + definition = """ + -> Master + --- + n : int32 + """ + + @property + def key_source(self): + return Master & Master.PartA & Master.PartB + + def make(self, key): + if key["master_id"] >= 4: + raise ValueError("fails on purpose") + self.insert1(dict(key, n=len(Master.PartA & key) * len(Master.PartB & key))) + + return Master, Summary + + +def _add_master(Master, master_id, n_parts=3): + Master.insert1(dict(master_id=master_id)) + Master.PartA.insert([dict(master_id=master_id, a=i) for i in range(n_parts)]) + Master.PartB.insert([dict(master_id=master_id, b=i) for i in range(n_parts)]) + + +def _master_statuses(jobs): + return {row["master_id"]: row["status"] for row in jobs.proj("status").to_dicts()} + + +def _populate_and_vanish(Master, Summary): + """Six masters: 0-3 succeed, 4-5 fail; then the keys of 0, 1 and 4 leave key_source.""" + for m in range(6): + _add_master(Master, m) + Summary.populate(reserve_jobs=True, suppress_errors=True) + assert _master_statuses(Summary.jobs) == { + 0: "success", + 1: "success", + 2: "success", + 3: "success", + 4: "error", + 5: "error", + } + (Master & {"master_id": 0}).delete(prompt=False) # key and target row go + (Master.PartA & {"master_id": 1}).delete_quick() # key leaves key_source; the target row stays + (Master.PartB & {"master_id": 4}).delete_quick() # an error whose key left key_source + + +def test_refresh_master_with_two_parts(one_master): + """Vanished success and error records are removed; the others are kept.""" + Master, Summary = one_master + with dj.config.override(jobs={"keep_completed": True}): + _populate_and_vanish(Master, Summary) + + # 3a: success/error records whose key is in neither key_source nor the target + result = Summary.jobs.refresh() + assert result["removed"] == 2 + assert _master_statuses(Summary.jobs) == {1: "success", 2: "success", 3: "success", 5: "error"} + + # 3b: the success record of master 1 still has a target row; it goes once stale + result = Summary.jobs.refresh(stale_timeout=0) + assert result["removed"] == 1 + assert _master_statuses(Summary.jobs) == {2: "success", 3: "success", 5: "error"} + + # nothing left to remove + assert Summary.jobs.refresh(stale_timeout=0)["removed"] == 0 + + +def _negated_subqueries(sql): + """The text of every ``NOT IN (...)`` subquery in ``sql``.""" + found = [] + lowered = sql.lower() + pos = lowered.find("not in (") + while pos >= 0: + start = pos + len("not in (") + depth, end = 1, start + while depth and end < len(sql): + depth += {"(": 1, ")": -1}.get(sql[end], 0) + end += 1 + found.append(lowered[start:end]) + pos = lowered.find("not in (", start) + return found + + +def test_refresh_never_negates_key_source(one_master, monkeypatch): + """No refresh statement nests the part-restricted key_source inside NOT IN.""" + Master, Summary = one_master + parts = [Master.PartA.table_name.lower(), Master.PartB.table_name.lower()] + with dj.config.override(jobs={"keep_completed": True}): + _populate_and_vanish(Master, Summary) + jobs = Summary.jobs + queries = [] + query = jobs.connection.query + + def recording_query(sql, *args, **kwargs): + queries.append(sql) + return query(sql, *args, **kwargs) + + monkeypatch.setattr(jobs.connection, "query", recording_query) + removed = jobs.refresh(stale_timeout=0)["removed"] + monkeypatch.undo() + + assert removed == 3 + offending = [sql for sql in queries if any(p in sub for sub in _negated_subqueries(sql) for p in parts)] + assert not offending, offending + + +def test_refresh_pages_through_candidates(one_master, monkeypatch): + """With more candidates than one page, every vanished job is found and the rest kept.""" + Master, Summary = one_master + monkeypatch.setattr(dj.jobs, "VANISHED_DELETE_BATCH", 2) + with dj.config.override(jobs={"keep_completed": True}): + for m in range(4): + _add_master(Master, m) + Summary.populate(reserve_jobs=True) + for m in range(4, 9): # jobs that make() never ran: pending + _add_master(Master, m) + Summary.jobs.refresh() + assert len(Summary.jobs) == 9 + # vanish keys on page boundaries and inside pages + for m in (0, 1, 4, 6, 8): + (Master.PartB & {"master_id": m}).delete_quick() + + result = Summary.jobs.refresh(stale_timeout=0) + assert result["removed"] == 5 + assert _master_statuses(Summary.jobs) == {2: "success", 3: "success", 5: "pending", 7: "pending"} + + +def test_refresh_moves_past_keys_that_do_not_match_themselves(schema_by_backend, monkeypatch): + """A page whose keys cannot be matched by value does not stop the cleanup. + + A float32 key such as 0.1 is stored inexactly, so restricting by its fetched + value finds no row. A page of such keys is neither found present nor deleted; + the cleanup must still move on to the next page instead of reading the same + page again. + """ + + @schema_by_backend + class Source(dj.Manual): + definition = """ + x : float32 + """ + + @schema_by_backend + class Target(dj.Computed): + definition = """ + -> Source + """ + + def make(self, key): + self.insert1(key) + + Source.insert([{"x": x} for x in (0.1, 0.2, 0.3, 0.5, 1.0, 2.0)]) + Target.jobs.refresh() + assert len(Target.jobs) == 6 + + monkeypatch.setattr(dj.jobs, "VANISHED_DELETE_BATCH", 3) + calls = 0 + keys = dj.expression.QueryExpression.keys + + def counted_keys(self, *args, **kwargs): + nonlocal calls + calls += 1 + assert calls < 50, "refresh() keeps reading the same page" + return keys(self, *args, **kwargs) + + monkeypatch.setattr(dj.expression.QueryExpression, "keys", counted_keys) + (Source & {"x": 2.0}).delete_quick() + + result = Target.jobs.refresh(stale_timeout=0) + + # The first page (0.1, 0.2, 0.3) cannot be matched and is skipped; on the + # second page, 2.0 has vanished and is removed. + assert result["removed"] == 1 + assert len(Target.jobs) == 5