diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index fa84471d36..d1bea819eb 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -268,6 +268,7 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_fix_sentencepiece_tokenizer_guard.py \ tests/saving/test_compressed_export_schemes.py \ tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ @@ -358,6 +359,7 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_fix_sentencepiece_tokenizer_guard.py \ tests/saving/test_compressed_export_schemes.py \ tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py index 457dfc8ea2..0cd702bce2 100644 --- a/studio/backend/tests/test_training_stop_watchdog.py +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -258,15 +258,27 @@ def test_watchdog_no_op_when_worker_superseded(monkeypatch): def test_new_run_gets_its_own_watchdog(monkeypatch): # A stale watchdog sleeping on an old proc must not stop a new run's stop from # creating its own watcher. - monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) - monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) b = TrainingBackend() - _record_force_terminate(monkeypatch, b) + started = [] + release = threading.Event() + + def _blocked_watchdog( + target_proc, + cancel, + watched_job_id = None, + ): + started.append(target_proc) + # No timeout: the finally always releases this, so a superseded watchdog stays + # alive through the assertions regardless of load; as a daemon it can't hang exit. + release.wait() + + monkeypatch.setattr(b, "_stop_watchdog_loop", _blocked_watchdog) old_proc = _FakeProc(alive = True) b._proc = old_proc b._start_stop_watchdog(cancel = False) first_wd = b._stop_watchdog + assert _wait_until(lambda: started == [old_proc]) # New run: fresh worker replaces the handle; its stop must get a new watcher # even though the old (superseded) watchdog is still alive. @@ -276,12 +288,12 @@ def test_new_run_gets_its_own_watchdog(monkeypatch): second_wd = b._stop_watchdog try: + assert _wait_until(lambda: started == [old_proc, new_proc]) assert first_wd.is_alive() assert second_wd is not first_wd, "a new run must get its own watchdog" assert b._stop_watchdog_proc is new_proc finally: - old_proc._alive = False - new_proc._alive = False + release.set() first_wd.join(timeout = 5) second_wd.join(timeout = 5) diff --git a/tests/saving/test_fix_sentencepiece_tokenizer_guard.py b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py new file mode 100644 index 0000000000..1ee523d57b --- /dev/null +++ b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: AGPL-3.0-only +import gc +import os + +os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python") + +import transformers +from transformers.utils import sentencepiece_model_pb2 + +from unsloth.tokenizer_utils import fix_sentencepiece_tokenizer + + +NORMAL, CONTROL = 1, 3 + + +def _spm_bytes(pieces): + m = sentencepiece_model_pb2.ModelProto() + for piece, score, typ in pieces: + p = m.pieces.add() + p.piece = piece + p.score = score + p.type = typ + return m.SerializeToString() + + +def _read_pieces(path): + m = sentencepiece_model_pb2.ModelProto() + with open(path, "rb") as f: + m.ParseFromString(f.read()) + return [p.piece for p in m.pieces] + + +class _FakeTokenizer: + """Minimal stand-in for a sentencepiece-backed slow tokenizer. + + ``save_pretrained`` writes a tokenizer.model, which is what the real slow + tokenizers do and what fix_sentencepiece_tokenizer reads back. + """ + + def __init__( + self, + name, + spm_bytes = None, + vocab = None, + ): + self.name = name + self.eos_token = "" + self.pad_token = "" + self._spm_bytes = spm_bytes + self._vocab = vocab or {} + self.saved_to = [] + + def save_pretrained(self, location): + self.saved_to.append(location) + os.makedirs(location, exist_ok = True) + if self._spm_bytes is not None: + with open(os.path.join(location, "tokenizer.model"), "wb") as f: + f.write(self._spm_bytes) + + def __call__( + self, + texts, + add_special_tokens = False, + ): + class _Encoded: + pass + + encoded = _Encoded() + encoded.input_ids = [[self._vocab[text]] for text in texts] + return encoded + + +def _tokenizers(): + pieces = [("", 0.0, CONTROL), ("a", -1.0, NORMAL), ("", 0.0, CONTROL)] + old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2}) + new = _FakeTokenizer("new") + return old, new + + +class _ReloadedTokenizer: + """Weakref-able stand-in for the tokenizer AutoTokenizer.from_pretrained returns.""" + + def __init__(self, location): + self.location = location + + +def _stub_auto_tokenizer(monkeypatch): + """fix_sentencepiece_tokenizer reloads the patched directory through + AutoTokenizer at the end; that needs a full tokenizer on disk, which is + out of scope here. Record the reload location and hand back a sentinel. + """ + loaded = [] + + class _StubAutoTokenizer: + @staticmethod + def from_pretrained(location, **kwargs): + loaded.append(location) + return _ReloadedTokenizer(location) + + monkeypatch.setattr(transformers, "AutoTokenizer", _StubAutoTokenizer) + return loaded + + +def test_old_tokenizer_is_saved_so_its_model_can_be_read(tmp_path, monkeypatch): + """The guard must not skip the body on a fresh temporary directory. + + fix_sentencepiece_tokenizer creates its scratch directory itself and then + checks for a tokenizer.model inside it, but that file only appears once + old_tokenizer.save_pretrained() has run. + """ + _stub_auto_tokenizer(monkeypatch) + old, new = _tokenizers() + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + + assert old.saved_to, "old tokenizer was never saved: the body did not run" + + +def test_token_mapping_is_applied_to_the_sentencepiece_model(tmp_path, monkeypatch): + loaded = _stub_auto_tokenizer(monkeypatch) + old, new = _tokenizers() + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + # Hold the returned tokenizer so its scratch dir survives until we read it. + tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + + assert "<|im_end|>" in _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert tok is not None + + +def test_tokenizer_without_a_sentencepiece_model_is_returned_untouched(tmp_path, monkeypatch): + """A fast-only tokenizer writes no tokenizer.model, so the guard still + short-circuits and the caller gets new_tokenizer back unchanged. Its scratch + dir is unreferenced and reclaimed immediately. + """ + _stub_auto_tokenizer(monkeypatch) + old = _FakeTokenizer("old", spm_bytes = None) + new = _FakeTokenizer("new") + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + result = fix_sentencepiece_tokenizer( + old, new, {"": "<|im_end|>"}, temporary_location = location + ) + + assert result is new + assert not any( + name.startswith("tokenizer_") for name in os.listdir(location) + ), "the fast-only scratch dir was not reclaimed" + + +def test_each_call_uses_a_fresh_isolated_subdirectory(tmp_path, monkeypatch): + """Each call must work in its own unique subdirectory, so concurrent or + repeated calls never share scratch files, stale artifacts never leak into + the reload, and nothing the caller left in the scratch location is deleted. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + os.makedirs(location, exist_ok = True) + + # A pre-existing artifact in the shared scratch location. + marker = os.path.join(location, "leftover.json") + with open(marker, "w") as f: + f.write("{}") + + old1, new1 = _tokenizers() + old2, new2 = _tokenizers() + # Hold both returned tokenizers so their scratch dirs stay alive. + tok1 = fix_sentencepiece_tokenizer( + old1, new1, {"": "<|im_end|>"}, temporary_location = location + ) + tok2 = fix_sentencepiece_tokenizer( + old2, new2, {"": "<|im_end|>"}, temporary_location = location + ) + + work1, work2 = loaded[0], loaded[1] + assert work1 != work2, "two calls reused the same directory" + assert os.path.dirname(work1) == location and os.path.dirname(work2) == location + assert os.path.isdir(work1) and os.path.isdir(work2) + # Nothing the caller left behind is deleted, and it never leaks into a work dir. + assert os.path.isfile(marker), "a pre-existing scratch file was deleted" + assert not os.path.isfile(os.path.join(work1, "leftover.json")) + assert not os.path.isfile(os.path.join(work2, "leftover.json")) + assert tok1 is not None and tok2 is not None + + +def test_sentencepiece_scratch_dir_is_reclaimed_once_the_tokenizer_is_gone(tmp_path, monkeypatch): + """The scratch dir must live as long as the returned tokenizer (its vocab_file + points there), then be reclaimed when the tokenizer is garbage collected. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + old, new = _tokenizers() + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + work = loaded[-1] + assert os.path.isdir(work), "scratch dir vanished while the tokenizer was alive" + + del tok + gc.collect() + assert not os.path.isdir(work), "scratch dir was not reclaimed after the tokenizer was freed" + + +class _CopyFromSubdirTokenizer: + """A slow tokenizer whose sentencepiece source lives elsewhere (like the + tokenizers convert_to_fast_tokenizer produces under {location}/{name}). + save_pretrained copies that source into the destination, as HF slow + tokenizers copy their vocab_file. + """ + + def __init__(self, source_model_path): + self.eos_token = "" + self.pad_token = "" + self._source_model_path = source_model_path + + def save_pretrained(self, location): + os.makedirs(location, exist_ok = True) + if os.path.isfile(self._source_model_path): + with open(self._source_model_path, "rb") as src: + data = src.read() + with open(os.path.join(location, "tokenizer.model"), "wb") as dst: + dst.write(data) + + def __call__( + self, + texts, + add_special_tokens = False, + ): + class _Encoded: + pass + + encoded = _Encoded() + encoded.input_ids = [[2] for _ in texts] + return encoded + + +def test_source_vocab_outside_the_work_directory_is_not_disturbed(tmp_path, monkeypatch): + """A tokenizer whose sentencepiece source lives elsewhere (e.g. the subtree + convert_to_fast_tokenizer created) is copied into the fresh work directory + and patched there; the original source is left untouched. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + subdir = os.path.join(location, "some_model") + os.makedirs(subdir, exist_ok = True) + + pieces = [("", 0.0, CONTROL), ("a", -1.0, NORMAL), ("", 0.0, CONTROL)] + source_model = os.path.join(subdir, "tokenizer.model") + with open(source_model, "wb") as f: + f.write(_spm_bytes(pieces)) + + old = _CopyFromSubdirTokenizer(source_model) + new = _FakeTokenizer("new") + tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + + assert _read_pieces(source_model) == [ + "", + "a", + "", + ], "the original source vocab was modified" + assert "<|im_end|>" in _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert tok is not None + + +def test_swap_mapping_swaps_both_pieces_without_duplicating(tmp_path, monkeypatch): + """When the caller swaps eos and stop_word in the fast JSON it must pass both + directions here; a one-way mapping would leave two stop_word pieces and no eos. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + pieces = [("", 0.0, CONTROL), ("<|im_end|>", -1.0, NORMAL), ("", 0.0, CONTROL)] + old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2, "<|im_end|>": 1}) + new = _FakeTokenizer("new") + + tok = fix_sentencepiece_tokenizer( + old, new, {"": "<|im_end|>", "<|im_end|>": ""}, temporary_location = location + ) + + result = _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert result.count("<|im_end|>") == 1 and result.count("") == 1, result + assert tok is not None + + +def test_only_applied_mappings_are_patched(tmp_path, monkeypatch): + """When the caller skips a mapping whose target already exists, it must not + pass that mapping here, or the skipped source token gets renamed anyway and + duplicates the existing target in the model. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + pieces = [ + ("", 0.0, CONTROL), + ("aa", -1.0, NORMAL), + ("bb", -1.0, NORMAL), + ("X", -1.0, NORMAL), + ] + old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"aa": 1, "bb": 2}) + new = _FakeTokenizer("new") + + # Caller skipped aa->X (X already exists) and applied bb->Y, so only bb->Y is passed. + tok = fix_sentencepiece_tokenizer(old, new, {"bb": "Y"}, temporary_location = location) + + result = _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert result.count("X") == 1 and "Y" in result and "aa" in result, result + assert tok is not None diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py index 7cad6cacfe..73d1244ee3 100644 --- a/tests/studio/test_studio_text_descender_clipping.py +++ b/tests/studio/test_studio_text_descender_clipping.py @@ -34,17 +34,22 @@ def test_model_selector_trigger_label_uses_leading_tight(): def test_sidebar_account_block_uses_leading_tight(): src = _read(APP_SIDEBAR) - # Match the account-block parent div regardless of its gap utility; this - # guard is about the leading-* class, not the spacing. - pattern = re.compile( - r'', - ) - matches = pattern.findall(src) + class_names = re.findall(r'