Upload 141 files

hugging/td_fuse/merge.py

@@ -386,7 +386,36 @@ def save_checkpoint(
     cfg: MergeConfig,
 ):
     """Save a checkpoint after a successful merge stage."""
-    ckpt_dir = Path(cfg.checkpoint_dir) / f"after_{stage_name}"
+    import shutil
+
+    ckpt_base = Path(cfg.checkpoint_dir)
+    ckpt_dir = ckpt_base / f"after_{stage_name}"
+
+    # --- Pre-save cleanup: free disk space ---
+    # 1. Delete residuals (non-essential, 5-20GB)
+    residuals_dir = ckpt_base / "residuals"
+    if residuals_dir.exists():
+        shutil.rmtree(str(residuals_dir), ignore_errors=True)
+        print(f"[merge] Freed disk: deleted residuals")
+
+    # 2. Delete td_fuse_outputs/final (duplicate of last checkpoint, ~17GB)
+    final_dir = Path("td_fuse_outputs") / "final"
+    if final_dir.exists():
+        shutil.rmtree(str(final_dir), ignore_errors=True)
+        print(f"[merge] Freed disk: deleted td_fuse_outputs/final")
+
+    # 3. Delete OLD checkpoints (already on HuggingFace via watcher)
+    if ckpt_base.exists():
+        for old_ckpt in ckpt_base.glob("after_*"):
+            if old_ckpt.name != f"after_{stage_name}" and old_ckpt.is_dir():
+                shutil.rmtree(str(old_ckpt), ignore_errors=True)
+                print(f"[merge] Freed disk: deleted old checkpoint {old_ckpt.name}")
+
+    # Check disk space
+    import shutil as sh_util
+    total, used, free = sh_util.disk_usage("/")
+    print(f"[merge] Disk after cleanup: {free/1e9:.1f} GB free / {total/1e9:.1f} GB total")
+
     ckpt_dir.mkdir(parents=True, exist_ok=True)
 
     print(f"[merge] Saving checkpoint to {ckpt_dir}...")
@@ -1158,24 +1187,18 @@ def run_pipeline(
         healing_targets = residual_bank.get_healing_targets(top_n=50)
         pipeline_results["suggested_healing_targets"] = healing_targets
 
-    # --- Save final model ---
+    # --- Skip final model save (duplicate of checkpoint, wastes 17GB disk) ---
+    # The checkpoint in td_fuse_checkpoints/after_<stage> IS the final model
     if pipeline_results["final_checkpoint"]:
-        final_dir = Path(cfg.output_dir) / "final"
-        final_dir.mkdir(parents=True, exist_ok=True)
-        # Free disk space before final save (Bug #25 fix)
+        pipeline_results["final_model_path"] = pipeline_results["final_checkpoint"]
+        print(f"\n[pipeline] Final model is at: {pipeline_results['final_checkpoint']}")
+        # Clean up models/base if still around
         import shutil as _shutil
-        for _cleanup in ["models/base"]:
+        for _cleanup in ["models/base", "td_fuse_outputs/final"]:
             _cp = Path(_cleanup)
             if _cp.exists() and _cp.is_dir():
                 _shutil.rmtree(str(_cp))
                 print(f"[merge] Freed disk: {_cleanup}")
-        import gc; gc.collect()
-        _stat = _shutil.disk_usage("/")
-        print(f"[merge] Disk: {_stat.free / 1e9:.1f} GB free / {_stat.total / 1e9:.1f} GB total")
-        target_model.save_pretrained(final_dir)
-        target_tokenizer.save_pretrained(final_dir)
-        pipeline_results["final_model_path"] = str(final_dir)
-        print(f"\n[pipeline] Final model saved to {final_dir}")
 
     if all_passed:
         pipeline_results["overall_status"] = "all_passed"