Upload 133 files

app.py

@@ -296,8 +296,8 @@ def _recover_after_obliterate():
                 _ts = datetime.now().strftime("%H:%M")
                 _short = model_choice.split("/")[-1] if "/" in model_choice else model_choice
                 _label = f"{method} on {_short} ({_ts}) [recovered]"
-                _last_obliterated_label = _label
                 with _lock:
+                    _last_obliterated_label = _label
                     _session_models[_label] = {
                         "model_id": data.get("model_id", model_choice),
                         "model_choice": model_choice,
@@ -328,9 +328,11 @@ def _recover_after_obliterate():
             f"was saved before the timeout. Switch to the **Chat** tab to use it. "
             f"Verification metrics were skipped."
         )
+        with _lock:
+            _label_snap = _last_obliterated_label
         dd = gr.update(
             choices=_get_session_model_choices(),
-            value=_last_obliterated_label or None,
+            value=_label_snap or None,
         )
         return status_msg, log_text, get_chat_header(), dd, gr.update(), dd
     else:
@@ -442,16 +444,17 @@ def _recover_sessions_from_disk() -> None:
                     "source": data.get("source", "recovered"),
                 }
             found_any = True
-            # Track the latest for auto-select
-            _last_obliterated_label = label
-            # Keep counter above any existing numbered dirs
-            if p.name.startswith("obliterated_"):
-                try:
-                    idx = int(p.name.split("_", 1)[1])
-                    if idx >= _obliterate_counter:
-                        _obliterate_counter = idx + 1
-                except (ValueError, IndexError):
-                    pass
+            # Track the latest for auto-select and keep counter above existing dirs.
+            # Protect globals with _lock to avoid races with concurrent obliterate().
+            with _lock:
+                _last_obliterated_label = label
+                if p.name.startswith("obliterated_"):
+                    try:
+                        idx = int(p.name.split("_", 1)[1])
+                        if idx >= _obliterate_counter:
+                            _obliterate_counter = idx + 1
+                    except (ValueError, IndexError):
+                        pass
     # If we recovered sessions and _state has no valid output_dir, set it to
     # the most recent checkpoint so chat_respond can reload from disk.
     # Also overwrite a stale output_dir that points to a non-existent path.
@@ -2133,6 +2136,33 @@ def _gpu_run_picklable(pipeline, fn, *args, **kwargs):
         pipeline._on_log = saved_on_log
 
 
+def _gpu_run_with_retry(pipeline, fn, *args, max_retries=2, stage_label="", on_log=None, **kwargs):
+    """Run a GPU stage via ``_gpu_run_picklable`` with automatic retry on ZeroGPU abort.
+
+    ZeroGPU can transiently abort GPU tasks due to timeouts, concurrent user
+    conflicts, or infrastructure issues.  Retrying often succeeds.  This wrapper
+    retries up to *max_retries* times with exponential backoff (3s, 9s) before
+    re-raising the final error.
+    """
+    last_exc = None
+    for attempt in range(1 + max_retries):
+        try:
+            return _gpu_run_picklable(pipeline, fn, *args, **kwargs)
+        except Exception as e:
+            last_exc = e
+            if not _is_zerogpu_abort(e) or attempt >= max_retries:
+                raise
+            delay = 3 * (3 ** attempt)  # 3s, 9s
+            if on_log:
+                on_log(
+                    f"[staged] GPU task aborted on attempt {attempt + 1} "
+                    f"({stage_label}) — retrying in {delay}s "
+                    f"({max_retries - attempt} retries left)..."
+                )
+            time.sleep(delay)
+    raise last_exc  # unreachable, but satisfies type checkers
+
+
 def obliterate(model_choice: str, method_choice: str,
                prompt_volume_choice: str, dataset_source_choice: str,
                custom_harmful: str, custom_harmless: str,
@@ -2409,27 +2439,27 @@ def obliterate(model_choice: str, method_choice: str,
                     if method == "informed":
                         # Informed pipeline: SUMMON+PROBE | ANALYZE+DISTILL+EXCISE | VERIFY+REBIRTH
                         on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
-                        _gpu_run_picklable(pipeline, pipeline.run_stage_summon_probe, time.time())
+                        _gpu_run_with_retry(pipeline, pipeline.run_stage_summon_probe, time.time(), stage_label="Stage 1: SUMMON+PROBE", on_log=on_log)
                         on_log("[staged] GPU released after Stage 1\n")
 
                         on_log("\u26a1 [staged] GPU Stage 2/3: ANALYZE + DISTILL + EXCISE")
-                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_analyze_distill_excise")
+                        _gpu_run_with_retry(pipeline, _restore_and_run_stage, pipeline, "run_stage_analyze_distill_excise", stage_label="Stage 2: ANALYZE+DISTILL+EXCISE", on_log=on_log)
                         on_log("[staged] GPU released after Stage 2\n")
 
                         on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
-                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth_informed")
+                        _gpu_run_with_retry(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth_informed", stage_label="Stage 3: VERIFY+REBIRTH", on_log=on_log)
                     else:
                         # Standard pipeline: SUMMON+PROBE | DISTILL+EXCISE | VERIFY+REBIRTH
                         on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
-                        _gpu_run_picklable(pipeline, pipeline.run_stage_summon_probe, time.time())
+                        _gpu_run_with_retry(pipeline, pipeline.run_stage_summon_probe, time.time(), stage_label="Stage 1: SUMMON+PROBE", on_log=on_log)
                         on_log("[staged] GPU released after Stage 1\n")
 
                         on_log("\u26a1 [staged] GPU Stage 2/3: DISTILL + EXCISE")
-                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_distill_excise")
+                        _gpu_run_with_retry(pipeline, _restore_and_run_stage, pipeline, "run_stage_distill_excise", stage_label="Stage 2: DISTILL+EXCISE", on_log=on_log)
                         on_log("[staged] GPU released after Stage 2\n")
 
                         on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
-                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth")
+                        _gpu_run_with_retry(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth", stage_label="Stage 3: VERIFY+REBIRTH", on_log=on_log)
                 finally:
                     # Clean up staged state temp dir
                     import shutil as _shutil
@@ -3021,9 +3051,12 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
         if not checkpoint or not Path(checkpoint).exists():
             _recover_sessions_from_disk()
             checkpoint = _state.get("output_dir")
-        # If output_dir is still stale, scan session models for any valid checkpoint
+        # If output_dir is still stale, scan session models for any valid checkpoint.
+        # Snapshot values under lock to avoid RuntimeError from concurrent dict modification.
         if not checkpoint or not Path(checkpoint).exists():
-            for _sm in _session_models.values():
+            with _lock:
+                _sm_snapshot = list(_session_models.values())
+            for _sm in _sm_snapshot:
                 _sm_dir = _sm.get("output_dir")
                 if _sm_dir and Path(_sm_dir).exists():
                     checkpoint = _sm_dir
@@ -3221,12 +3254,19 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
                 and _state.get("tokenizer") is not None
             )
         if choice and _model_ok:
-            # Double-check model tensors aren't stale (meta device)
-            try:
-                _dev = next(_state["model"].parameters()).device
-                if _dev.type == "meta":
+            # Double-check model tensors aren't stale (meta device).
+            # Re-acquire lock to safely access model — it could become None
+            # between the first lock release and this check.
+            with _lock:
+                _model_ref = _state.get("model")
+            if _model_ref is not None:
+                try:
+                    _dev = next(_model_ref.parameters()).device
+                    if _dev.type == "meta":
+                        _model_ok = False
+                except Exception:
                     _model_ok = False
-            except Exception:
+            else:
                 _model_ok = False
         if choice and _model_ok:
             yield (
@@ -3250,24 +3290,27 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
 
         # If recovery didn't find the exact choice, check if model is loaded
         if choice not in _bench_configs:
+            # Read state under lock, but never yield while holding the lock —
+            # yield suspends the generator and would block all other threads.
             with _lock:
-                if _state["status"] == "ready" and _state["model"] is not None:
-                    yield (
-                        f"**Ready!** Model already loaded — just type in the chat below.",
-                        get_chat_header(),
-                    )
-                    return
-                # Check if we can reload from a checkpoint on disk
+                _is_ready = _state["status"] == "ready" and _state["model"] is not None
                 checkpoint = _state.get("output_dir")
-                if checkpoint and Path(checkpoint).exists():
-                    yield (
-                        f"**Loading model** from saved checkpoint...",
-                        "",
-                    )
+                _model_name_snap = _state.get("model_name") or ""
+            if _is_ready:
+                yield (
+                    f"**Ready!** Model already loaded — just type in the chat below.",
+                    get_chat_header(),
+                )
+                return
+            # Check if we can reload from a checkpoint on disk
+            if checkpoint and Path(checkpoint).exists():
+                yield (
+                    f"**Loading model** from saved checkpoint...",
+                    "",
+                )
             # If we have a checkpoint, attempt reload outside the lock
-            checkpoint = _state.get("output_dir")
             if checkpoint and Path(checkpoint).exists():
-                is_preset = (_state.get("model_name") or "") in MODELS
+                is_preset = _model_name_snap in MODELS
                 try:
                     model_loaded = _load_model_to_device(
                         checkpoint, torch_dtype=torch.float16,
@@ -3304,26 +3347,32 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
 
     # If this model is already the active one, skip the destructive reload
     with _lock:
-        if (_state["status"] == "ready"
-                and _state["model"] is not None
-                and _state["model_name"] == cfg.get("model_choice", "")
-                and _state["method"] == method_key):
-            yield (
-                f"**Already loaded!** `{choice}` is ready — just type in the chat below.",
-                get_chat_header(),
-            )
-            return
+        _already_active = (
+            _state["status"] == "ready"
+            and _state["model"] is not None
+            and _state["model_name"] == cfg.get("model_choice", "")
+            and _state["method"] == method_key
+        )
+    if _already_active:
+        yield (
+            f"**Already loaded!** `{choice}` is ready — just type in the chat below.",
+            get_chat_header(),
+        )
+        return
 
     # Unstick stale "obliterating" status left behind by ZeroGPU timeout
     _unstick_stale_obliterating()
 
     with _lock:
-        if _state["status"] == "obliterating":
-            yield "**Error:** An obliteration is already in progress.", ""
-            return
-        _state["status"] = "obliterating"
-        _state["model_name"] = cfg["model_choice"]
-        _state["method"] = method_key
+        _already_obliterating = _state["status"] == "obliterating"
+        if not _already_obliterating:
+            _state["status"] = "obliterating"
+            _state["obliterate_started_at"] = time.time()
+            _state["model_name"] = cfg["model_choice"]
+            _state["method"] = method_key
+    if _already_obliterating:
+        yield "**Error:** An obliteration is already in progress.", ""
+        return
     _clear_gpu()
 
     # If we have a saved checkpoint on disk, load directly — no re-training!
@@ -3348,6 +3397,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
                 _state["tokenizer"] = tokenizer_loaded
                 _state["steering"] = None
                 _state["status"] = "ready"
+                _state["obliterate_started_at"] = None
                 _state["output_dir"] = checkpoint_dir
             progress(1.0, desc="Ready!")
             yield (
@@ -3383,6 +3433,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
                     _state["tokenizer"] = tokenizer_loaded
                     _state["steering"] = None
                     _state["status"] = "ready"
+                    _state["obliterate_started_at"] = None
                     _state["output_dir"] = checkpoint_dir
                 progress(1.0, desc="Ready!")
                 yield (