fix(pipeline): fail closed on invalid output

This commit is contained in:
Joseph Magly
2026-08-24 13:42:13 -04:00
parent 7162203da1
commit d8d0a76231
7 changed files with 515 additions and 47 deletions
+97
View File
@@ -72,3 +72,100 @@ assert app._state["output_dir"] is None
)
assert result.returncode == 0, result.stdout + result.stderr
@pytest.mark.operator_ui
def test_pipeline_cancellation_and_local_checkpoint_reload_contracts(tmp_path):
"""Cover timeout grace outcomes and local-only reload classification."""
script = r'''
import pathlib
import sys
import threading
import time
from unittest.mock import Mock
import app
root = pathlib.Path(sys.argv[1])
checkpoint = root / "completed"
checkpoint.mkdir()
(checkpoint / "abliteration_metadata.json").write_text("{}", encoding="utf-8")
(checkpoint / "config.json").write_text("{}", encoding="utf-8")
(checkpoint / "tokenizer_config.json").write_text("{}", encoding="utf-8")
(checkpoint / "model.safetensors").write_bytes(b"weights")
assert app._resolve_local_checkpoint(checkpoint) == checkpoint.resolve()
calls = []
model = Mock()
model.to.return_value = model
tokenizer = Mock(pad_token=None, eos_token="<eos>")
app.AutoModelForCausalLM.from_pretrained = Mock(
side_effect=lambda source, **kwargs: (calls.append((source, kwargs)), model)[1]
)
app.AutoTokenizer.from_pretrained = Mock(return_value=tokenizer)
app.dev.supports_device_map_auto = lambda: True
app._load_model_to_device(checkpoint, local_files_only=True)
assert calls[0][0] == checkpoint
assert calls[0][1]["local_files_only"] is True
quantization = object()
loaded_model, loaded_tokenizer = app._reload_local_checkpoint(
checkpoint,
trust_remote_code=True,
model_kwargs={"quantization_config": quantization},
)
assert loaded_model is model and loaded_tokenizer is tokenizer
assert calls[-1][1]["quantization_config"] is quantization
assert calls[-1][1]["local_files_only"] is True
assert tokenizer.pad_token == "<eos>"
app._reload_local_checkpoint(
checkpoint,
trust_remote_code=True,
model_kwargs={"offload_folder": str(root / "offload")},
)
assert calls[-1][1]["offload_folder"] == str(root / "offload")
assert calls[-1][1]["local_files_only"] is True
assert app._format_checkpoint_reload_error(FileNotFoundError("gone")).startswith(
"Saved checkpoint is missing"
)
assert app._format_checkpoint_reload_error(ValueError("Checkpoint has no model weights")).startswith(
"Saved checkpoint is incomplete"
)
assert app._format_checkpoint_reload_error(ValueError("Checkpoint metadata is corrupt")).startswith(
"Saved checkpoint is corrupt"
)
assert app._format_checkpoint_reload_error(ValueError("Repo id is invalid")).startswith(
"Checkpoint Hub resolution failed"
)
def run_cooperative(cancel, delay):
cancel.wait()
time.sleep(delay)
cancel = threading.Event()
worker = threading.Thread(target=run_cooperative, args=(cancel, 0.02))
worker.start()
assert app._cancel_pipeline_worker(worker, cancel, grace_seconds=1.0) is True
assert cancel.is_set()
cancel = threading.Event()
release = threading.Event()
worker = threading.Thread(target=lambda: release.wait())
worker.start()
assert app._cancel_pipeline_worker(worker, cancel, grace_seconds=0.01) is False
assert cancel.is_set() and worker.is_alive()
release.set()
worker.join(timeout=1.0)
assert not worker.is_alive()
'''
result = subprocess.run(
[sys.executable, "-c", script, str(tmp_path)],
capture_output=True,
text=True,
timeout=120,
check=False,
)
assert result.returncode == 0, result.stdout + result.stderr