diff --git a/server.py b/server.py index 8678717..dd6cf58 100644 --- a/server.py +++ b/server.py @@ -84,8 +84,30 @@ def _parse_cache_max_completed_jobs(value: str | None) -> int: )) MAX_CONCURRENT = 2 CACHE_TTL = os.environ.get("CACHE_TTL", "24h") -# Directories under CACHE_DIR that hold model weights (never auto-deleted) -_PRESERVED_CACHE_DIRS = frozenset({"torch", "huggingface", "locale"}) +# Directories under CACHE_DIR that hold model weights (never auto-deleted). +# +# This list was a BLACKLIST, and a blacklist is the wrong shape for this job: it has to +# enumerate everything that must survive, so anything it forgets gets deleted. It forgot +# `_roformer-models`, and the 24h sweeper duly deleted the 700 MB BS-Roformer-SW checkpoint +# every single day — users re-downloaded it on the next start, forever, with no error and no +# clue why. (Reported at got-feedBack/feedBack-plugin-stem-splitter, seen in the wild.) +# +# The real fix is _CACHE_ENTRY_RE below: only delete directories that LOOK like stem-cache +# entries. The list stays as a second line of defence. +_PRESERVED_CACHE_DIRS = frozenset({"torch", "huggingface", "locale", "_roformer-models"}) + +# What a stem-cache directory is named: _job_id_for() -> f"{audio_hash}-{slug}", where the +# hash is sha256[:16] (hex) and the slug is a sanitized model name. +# +# The sweeper only deletes things matching THIS. Anything else under CACHE_DIR — a model dir +# we added, a model dir some future version adds, a stray file a user dropped in — is left +# alone. Deleting only what we recognize is the only version of this that stays correct as +# the cache dir grows new neighbours. +# Exactly what _job_id_for() emits: sha256[:16] (lowercase hex) + "-" + a slug that is +# re.sub(r"[^A-Za-z0-9]+", "-", model).strip("-").lower() — so lowercase, digits and +# hyphens, nothing else. Matching more loosely than we emit only widens the set of +# directories we are willing to delete, which is the opposite of the point. +_CACHE_ENTRY_RE = re.compile(r"^[0-9a-f]{16}-[a-z0-9-]+$") def _parse_ttl(ttl_str: str) -> int | None: @@ -2193,9 +2215,19 @@ def _run_warmup() -> None: def _cache_cleanup_loop() -> None: """Background daemon thread: periodically delete expired stem cache dirs. - - Walks CACHE_DIR, skips preserved directories (torch, huggingface, locale), - and deletes any stem cache directory whose mtime exceeds CACHE_TTL. + + Walks CACHE_DIR and deletes a directory only if it BOTH: + + * is not in _PRESERVED_CACHE_DIRS (torch, huggingface, locale, _roformer-models), and + * matches _CACHE_ENTRY_RE — i.e. it looks like something _job_id_for() produced. + + The second condition is the load-bearing one. This used to be a preserve-list alone, which + is a blacklist: it has to enumerate everything that must survive, so whatever it forgets + gets destroyed. It forgot `_roformer-models`, and duly deleted the 700 MB BS-Roformer-SW + checkpoint every 24 hours — users re-downloaded it on the next start, forever, with no + error and no clue why. Deleting only what we RECOGNIZE is the only version of this that + stays correct as the cache dir gains new neighbours. + Runs every 10 minutes. """ if CACHE_TTL_SECONDS is None: @@ -2212,6 +2244,10 @@ def _cache_cleanup_loop() -> None: continue if entry.name in _PRESERVED_CACHE_DIRS: continue + if not _CACHE_ENTRY_RE.fullmatch(entry.name): + # Not a stem-cache entry -> not ours to delete. Model weights, and + # anything else that comes to live under CACHE_DIR, land here. + continue try: mtime = entry.stat().st_mtime age = now - mtime diff --git a/tests/test_cache_cleanup.py b/tests/test_cache_cleanup.py index 1dbfd96..e995ee8 100644 --- a/tests/test_cache_cleanup.py +++ b/tests/test_cache_cleanup.py @@ -276,3 +276,96 @@ def test_sleep_at_end_of_loop_not_beginning(self): f"time.sleep should be the LAST statement in the while loop, " f"but last statement is: {last_stmt_text[:100]}" ) + +# ── the sweeper must not eat model weights ────────────────────────────────── + +class TestSweeperOnlyDeletesStemCaches: + """The 24h TTL sweeper deleted the 700 MB BS-Roformer-SW checkpoint every day. + + Model weights live under CACHE_DIR alongside the stem caches, and the sweeper protected + them with a BLACKLIST (`torch`, `huggingface`, `locale`). It forgot `_roformer-models`. + So every 24 hours the checkpoint was deleted, the next start re-downloaded 700 MB, and + nothing anywhere said why. Seen in the wild. + + A blacklist is the wrong shape here: it must enumerate everything that has to survive, so + whatever it forgets gets destroyed — including any model dir a future version adds. The + sweeper now deletes only names that LOOK like stem-cache entries. + + Everything below is read out of server.py with AST. Parsing the source (rather than + grepping it, or exec'ing a line of it) means these tests keep testing the CODE and not the + text: they don't pass because a name appears in a comment, and they don't break when a + line moves. + """ + + TREE = ast.parse(SERVER_PY.read_text(encoding="utf-8")) + + def _literal(self, name): + """The value assigned to a module-level constant.""" + for node in ast.iter_child_nodes(self.TREE): + if isinstance(node, ast.Assign): + for tgt in node.targets: + if isinstance(tgt, ast.Name) and tgt.id == name: + return node.value + raise AssertionError(f"{name} not found in server.py") + + def _entry_re(self): + # _CACHE_ENTRY_RE = re.compile(r"...") -> compile the literal pattern itself. + call = self._literal("_CACHE_ENTRY_RE") + # Assert it is specifically a re.compile(...) call, not merely "a call". Otherwise a + # refactor to some other callable with the same first-argument shape would slip + # through, and this helper would compile whatever that call's first argument happened + # to be — testing a pattern the sweeper does not use. + assert isinstance(call, ast.Call), "_CACHE_ENTRY_RE should be a re.compile(...) call" + assert ast.unparse(call.func) == "re.compile", ( + f"_CACHE_ENTRY_RE should be built by re.compile, got {ast.unparse(call.func)}") + pattern = ast.literal_eval(call.args[0]) + return re.compile(pattern) + + def _preserved(self): + # frozenset({...}) -> the real set, not a substring of the file. + node = self._literal("_PRESERVED_CACHE_DIRS") + if isinstance(node, ast.Call): # frozenset({...}) + node = node.args[0] + return set(ast.literal_eval(node)) + + # ── what may be deleted ─────────────────────────────────────────────────── + def test_a_real_stem_cache_entry_is_deletable(self): + # _job_id_for() -> f"{sha256[:16]}-{model-slug}" + assert self._entry_re().fullmatch("7a819518228bca22-bs-roformer-sw") + assert self._entry_re().fullmatch("deadbeefdeadbeef-htdemucs-ft") + + def test_the_roformer_model_dir_is_NOT_deletable(self): + """The actual bug. This directory holds a 700 MB checkpoint.""" + assert not self._entry_re().fullmatch("_roformer-models") + + def test_the_other_model_dirs_are_not_deletable(self): + for name in ("torch", "huggingface", "locale", "hub"): + assert not self._entry_re().fullmatch(name), name + + def test_an_unknown_future_model_dir_is_not_deletable(self): + """The point of a whitelist: a model dir nobody has added yet must survive too. Under + the old blacklist it would have been silently deleted after 24h.""" + for name in ("_mdx-models", "whisper-models", "some-new-cache", "openvino"): + assert not self._entry_re().fullmatch(name), name + + def test_the_pattern_is_no_looser_than_what_we_emit(self): + """Matching more loosely than _job_id_for() emits only widens the set of directories + we are willing to delete — the opposite of the point.""" + for name in ("DEADBEEFDEADBEEF-model", # hash is lowercase hex + "deadbeefdeadbeef-Model_X.1", # slug is lowercased + sanitized + "some.model.cache"): + assert not self._entry_re().fullmatch(name), name + + # ── the safeguards are real, not textual ────────────────────────────────── + def test_roformer_dir_is_also_on_the_preserve_list(self): + # Membership in the actual set — not "the string appears somewhere in the file", + # which would still pass if the name were only left behind in a comment. + assert "_roformer-models" in self._preserved() + + def test_the_sweeper_actually_consults_the_pattern(self): + """A constant nobody reads fixes nothing. Assert the cleanup loop REFERENCES it.""" + loop = next(n for n in ast.walk(self.TREE) + if isinstance(n, ast.FunctionDef) and n.name == "_cache_cleanup_loop") + names = {n.id for n in ast.walk(loop) if isinstance(n, ast.Name)} + assert "_CACHE_ENTRY_RE" in names, "_cache_cleanup_loop must check the pattern" + assert "_PRESERVED_CACHE_DIRS" in names