CalvinBackup/OBLITERATUS
1
0
Fork 0
mirror of https://github.com/elder-plinius/OBLITERATUS.git synced 2026-08-30 06:30:37 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(pipeline): fail closed on invalid output

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+178 -37
View File
@@ -71,6 +71,7 @@ from obliteratus.model_load_settings import (
ModelLoadSettings,
resolve_model_load_settings,
)
from obliteratus.persistence_contracts import validate_reloadable_checkpoint
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
# ── ZeroGPU support ─────────────────────────────────────────────────
@@ -121,6 +122,7 @@ def _load_model_to_device(
offload_folder: str | None = None,
low_cpu_mem_usage: bool = False,
token: str | None = None,
local_files_only: bool = False,
) -> AutoModelForCausalLM:
"""Load a causal LM onto the best available device, MPS-safe.
@@ -142,6 +144,8 @@ def _load_model_to_device(
kwargs["low_cpu_mem_usage"] = True
if token is not None:
kwargs["token"] = token
if local_files_only:
kwargs["local_files_only"] = True
if dev.supports_device_map_auto():
kwargs["device_map"] = "auto"
@@ -156,6 +160,87 @@ def _load_model_to_device(
return model
def _resolve_local_checkpoint(checkpoint: str | Path) -> Path:
"""Resolve and fully validate a local checkpoint before Transformers sees it."""
return validate_reloadable_checkpoint(checkpoint)
def _reload_local_checkpoint(
checkpoint: str | Path,
*,
trust_remote_code: bool,
model_kwargs: dict | None = None,
):
"""Load a validated local checkpoint without permitting Hub resolution."""
checkpoint_path = _resolve_local_checkpoint(checkpoint)
model = _load_model_to_device(
checkpoint_path,
trust_remote_code=trust_remote_code,
local_files_only=True,
**(model_kwargs or {}),
)
tokenizer = AutoTokenizer.from_pretrained(
checkpoint_path,
trust_remote_code=trust_remote_code,
local_files_only=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
def _format_checkpoint_reload_error(error: BaseException) -> str:
"""Return an actionable reload category without misreporting GPU capacity."""
message = str(error) or repr(error)
lowered = message.lower()
if isinstance(error, FileNotFoundError):
return f"Saved checkpoint is missing: {message}"
if "corrupt" in lowered:
return f"Saved checkpoint is corrupt: {message}"
if isinstance(error, ValueError) and "checkpoint" in lowered:
return f"Saved checkpoint is incomplete: {message}"
if dev.is_oom_error(error):
return f"Insufficient GPU memory to reload the completed checkpoint: {message}"
if "repo id" in lowered or "repository" in lowered:
return f"Checkpoint Hub resolution failed: {message}"
return f"Completed checkpoint could not be reloaded: {message}"
def _cancel_pipeline_worker(
worker: threading.Thread,
cancellation: threading.Event,
*,
grace_seconds: float,
) -> bool:
"""Signal cooperative cancellation and report confirmed worker termination."""
cancellation.set()
worker.join(timeout=grace_seconds)
return not worker.is_alive()
def _cleanup_failed_pipeline(pipeline, output_dir: str) -> None:
"""Release failed model state and remove only this UI run's final artifact."""
if pipeline is not None:
try:
pipeline.cleanup_failed_run()
except Exception:
pass
path = Path(output_dir)
try:
resolved = path.resolve()
temp_root = Path("/tmp").resolve()
if (
resolved.parent == temp_root
and resolved.name.startswith("obliterated_")
and resolved.is_dir()
):
import shutil
shutil.rmtree(resolved, ignore_errors=True)
except OSError:
pass
def _checkpoint_load_kwargs(load_settings: ModelLoadSettings) -> dict:
"""Translate resolved settings for direct loading of a saved checkpoint."""
torch_dtype = {
@@ -2072,6 +2157,7 @@ def obliterate(model_choice: str, method_choice: str,
last_yielded = [0]
pipeline_ref = [None]
error_ref = [None]
cancellation = threading.Event()
t_start = time.time()
def _elapsed():
@@ -2128,6 +2214,7 @@ def obliterate(model_choice: str, method_choice: str,
harmless_prompts=harmless_all[:n],
on_stage=on_stage,
on_log=on_log,
cancellation_event=cancellation,
)
pipeline_ref[0] = pipeline
pipeline.run_informed()
@@ -2145,6 +2232,7 @@ def obliterate(model_choice: str, method_choice: str,
harmless_prompts=harmless_all[:n],
on_stage=on_stage,
on_log=on_log,
cancellation_event=cancellation,
# Advanced overrides from UI
n_directions=int(adv_n_directions),
direction_method=adv_direction_method,
@@ -2210,6 +2298,7 @@ def obliterate(model_choice: str, method_choice: str,
# Stream log updates while pipeline runs (max 45 minutes to prevent indefinite hang)
_max_pipeline_secs = 45 * 60
_pipeline_start = time.time()
timed_out = False
status_msg = "**Obliterating\u2026** (0s)"
while worker.is_alive():
status_msg = f"**Obliterating\u2026** ({_elapsed()})"
@@ -2220,13 +2309,47 @@ def obliterate(model_choice: str, method_choice: str,
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
if time.time() - _pipeline_start > _max_pipeline_secs:
log_lines.append("\nTIMEOUT: Pipeline exceeded 45-minute limit.")
cancellation.set()
timed_out = True
break
time.sleep(0.5)
worker.join(timeout=30)
if timed_out:
terminated = _cancel_pipeline_worker(
worker,
cancellation,
grace_seconds=30,
)
pipeline = pipeline_ref[0]
if terminated:
_cleanup_failed_pipeline(pipeline, save_dir)
detail = "Worker stopped; incomplete output was discarded."
else:
detail = (
"Worker is still stopping cooperatively; no output or session state "
"will be exposed."
)
with _lock:
_state["status"] = "idle"
_state["model"] = None
_state["tokenizer"] = None
log_lines.append(detail)
_state["log"] = log_lines
yield (
"**Timed out:** Pipeline cancelled; no successful artifact was registered.",
"\n".join(log_lines),
get_chat_header(),
gr.update(),
gr.update(),
gr.update(),
)
return
worker.join()
# Handle error
if error_ref[0] is not None:
_cleanup_failed_pipeline(pipeline_ref[0], save_dir)
with _lock:
_state["status"] = "idle"
err_msg = str(error_ref[0]) or repr(error_ref[0])
@@ -2239,6 +2362,7 @@ def obliterate(model_choice: str, method_choice: str,
# Wrapped in try/except to ensure status is never stuck on "obliterating".
try:
pipeline = pipeline_ref[0]
save_dir = str(_resolve_local_checkpoint(save_dir))
can_generate = pipeline._quality_metrics.get("coherence") is not None
# ── Telemetry: log single obliteration to community leaderboard ──
@@ -2345,8 +2469,17 @@ def obliterate(model_choice: str, method_choice: str,
except ImportError:
pass
dtype_label = {
"float16": "FP16",
"bfloat16": "BF16",
"float32": "FP32",
}[load_settings.dtype]
checkpoint_path = _resolve_local_checkpoint(save_dir)
if bnb_available:
log_lines.append("\nModel too large for chat at float16 — reloading in 4-bit...")
log_lines.append(
f"\nModel too large for chat at {dtype_label} — reloading in 4-bit..."
)
last_yielded[0] = len(log_lines)
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
try:
@@ -2357,17 +2490,11 @@ def obliterate(model_choice: str, method_choice: str,
bnb_4bit_quant_type="nf4",
llm_int8_enable_fp32_cpu_offload=True,
)
model_reloaded = _load_model_to_device(
save_dir,
quantization_config=bnb_cfg,
model_reloaded, tokenizer_reloaded = _reload_local_checkpoint(
checkpoint_path,
trust_remote_code=True,
model_kwargs={"quantization_config": bnb_cfg},
)
tokenizer_reloaded = AutoTokenizer.from_pretrained(
save_dir,
trust_remote_code=True,
)
if tokenizer_reloaded.pad_token is None:
tokenizer_reloaded.pad_token = tokenizer_reloaded.eos_token
# Re-install activation steering hooks on the reloaded model
if steering_meta:
@@ -2382,14 +2509,17 @@ def obliterate(model_choice: str, method_choice: str,
can_generate = True
log_lines.append("Reloaded in 4-bit — chat is ready!")
except Exception as e:
log_lines.append(f"4-bit reload failed: {e}")
log_lines.append(
f"4-bit reload failed: {_format_checkpoint_reload_error(e)}"
)
_clear_gpu()
# -- Attempt 2: CPU offloading (slower but no extra dependencies)
if not can_generate:
import tempfile
log_lines.append(
"\nModel too large for chat at float16 — reloading with CPU offload..."
f"\nModel too large for chat at {dtype_label} — "
"reloading with CPU offload..."
if not bnb_available
else "Falling back to CPU offload..."
)
@@ -2397,22 +2527,18 @@ def obliterate(model_choice: str, method_choice: str,
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
try:
offload_dir = tempfile.mkdtemp(prefix="obliteratus_offload_")
model_reloaded = _load_model_to_device(
save_dir,
offload_folder=offload_dir,
torch_dtype={
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float32": torch.float32,
}[load_settings.dtype],
model_reloaded, tokenizer_reloaded = _reload_local_checkpoint(
checkpoint_path,
trust_remote_code=True,
model_kwargs={
"offload_folder": offload_dir,
"torch_dtype": {
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float32": torch.float32,
}[load_settings.dtype],
},
)
tokenizer_reloaded = AutoTokenizer.from_pretrained(
save_dir,
trust_remote_code=True,
)
if tokenizer_reloaded.pad_token is None:
tokenizer_reloaded.pad_token = tokenizer_reloaded.eos_token
# Re-install activation steering hooks on the reloaded model
if steering_meta:
@@ -2427,8 +2553,13 @@ def obliterate(model_choice: str, method_choice: str,
can_generate = True
log_lines.append("Reloaded with CPU offload — chat is ready (may be slower).")
except Exception as e:
log_lines.append(f"CPU offload reload failed: {e}")
log_lines.append("Chat unavailable. Load the saved model on a larger instance.")
log_lines.append(
f"CPU offload reload failed: {_format_checkpoint_reload_error(e)}"
)
log_lines.append(
"Chat unavailable. The completed checkpoint remains saved; "
"review the reload error above."
)
with _lock:
_state["status"] = "idle"
@@ -2474,6 +2605,7 @@ def obliterate(model_choice: str, method_choice: str,
except Exception as e:
# Ensure status never gets stuck on "obliterating"
_cleanup_failed_pipeline(pipeline_ref[0], save_dir)
with _lock:
_state["status"] = "idle"
err_msg = str(e) or repr(e)
@@ -2564,15 +2696,17 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
checkpoint = _state.get("output_dir")
if _checkpoint_is_available(checkpoint):
try:
checkpoint = _resolve_local_checkpoint(checkpoint)
is_preset = (_state.get("model_name") or "") in MODELS
load_settings = _settings_from_metadata(_state.get("load_settings"))
model = _load_model_to_device(
checkpoint,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer = AutoTokenizer.from_pretrained(
checkpoint, trust_remote_code=is_preset,
checkpoint, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
@@ -2584,8 +2718,8 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
_state["model"] = model
_state["tokenizer"] = tokenizer
_state["status"] = "ready"
except Exception:
yield "Model failed to reload from checkpoint. Try re-obliterating."
except Exception as error:
yield _format_checkpoint_reload_error(error)
return
else:
_clear_stale_model_state()
@@ -2778,14 +2912,16 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
if checkpoint and Path(checkpoint).exists():
is_preset = (_state.get("model_name") or "") in MODELS
try:
checkpoint = _resolve_local_checkpoint(checkpoint)
load_settings = _settings_from_metadata(_state.get("load_settings"))
model_loaded = _load_model_to_device(
checkpoint,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
checkpoint, trust_remote_code=is_preset,
checkpoint, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer_loaded.pad_token is None:
tokenizer_loaded.pad_token = tokenizer_loaded.eos_token
@@ -2799,7 +2935,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
)
return
except Exception as e:
yield f"**Error:** Could not reload model: {e}", get_chat_header()
yield f"**Error:** {_format_checkpoint_reload_error(e)}", get_chat_header()
return
yield (
"**Error:** Model checkpoint not found. The Space may have restarted — "
@@ -2847,13 +2983,15 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
is_preset = cfg["model_choice"] in MODELS
try:
checkpoint_dir = _resolve_local_checkpoint(checkpoint_dir)
model_loaded = _load_model_to_device(
checkpoint_dir,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
checkpoint_dir, trust_remote_code=is_preset,
checkpoint_dir, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer_loaded.pad_token is None:
tokenizer_loaded.pad_token = tokenizer_loaded.eos_token
@@ -2896,9 +3034,10 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
checkpoint_dir,
quantization_config=bnb_cfg,
trust_remote_code=is_preset,
local_files_only=True,
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
checkpoint_dir, trust_remote_code=is_preset,
checkpoint_dir, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer_loaded.pad_token is None:
tokenizer_loaded.pad_token = tokenizer_loaded.eos_token
@@ -3035,15 +3174,17 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[
model_name = _state.get("model_name") or model_name
if checkpoint and Path(checkpoint).exists():
try:
checkpoint = _resolve_local_checkpoint(checkpoint)
is_preset = (model_name or "") in MODELS
load_settings = _settings_from_metadata(_state.get("load_settings"))
abliterated_model = _load_model_to_device(
checkpoint,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer = AutoTokenizer.from_pretrained(
checkpoint, trust_remote_code=is_preset,
checkpoint, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token