Persist phase completion state to resume reliably

app.py

@@ -3,6 +3,8 @@
 import os
 import sys
 import shutil
+import json
+import time
 import gradio as gr
 import subprocess
 import threading
@@ -10,7 +12,7 @@ from pathlib import Path
 from datetime import datetime
 from typing import List, Tuple
 
-from huggingface_hub import login, hf_hub_download, HfApi
+from huggingface_hub import login, hf_hub_download, HfApi, create_repo
 from datasets import load_dataset, DatasetDict
 from src.config import (
     FORWARD_DATASET_NAME,
@@ -18,6 +20,7 @@ from src.config import (
     TOKENIZER_NAME,
     FORWARD_MODEL_NAME,
     RETRO_MODEL_NAME,
+    STATE_REPO,
 )
 
 # -----------------------------------------------------------------------------
@@ -37,6 +40,8 @@ FORWARD_MODEL_DIR = REPO_ROOT / "forward_model"
 RETRO_MODEL_DIR = REPO_ROOT / "retro_model"
 TOKENIZER_FILE = REPO_ROOT / "tokenizer.json"
 
+STATE_FILE = REPO_ROOT / "training_state.json"
+
 # Ensure working directories exist
 for path in (CACHE_DIR, HF_CACHE_DIR):
     path.mkdir(parents=True, exist_ok=True)
@@ -62,6 +67,65 @@ training_status = {
 HF_API = HfApi(token=HF_MODEL_TOKEN)
 WEIGHT_FILENAMES = {"pytorch_model.bin", "model.safetensors"}
 
+def load_training_state() -> dict:
+    if STATE_FILE.exists():
+        try:
+            with open(STATE_FILE, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception:
+            pass
+    if HF_MODEL_TOKEN:
+        try:
+            downloaded = hf_hub_download(
+                repo_id=STATE_REPO,
+                filename="training_state.json",
+                repo_type="dataset",
+                token=HF_MODEL_TOKEN,
+            )
+            shutil.copy(downloaded, STATE_FILE)
+            with open(STATE_FILE, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception:
+            return {}
+    return {}
+
+
+def save_training_state(state: dict):
+    if not HF_MODEL_TOKEN:
+        return
+    STATE_FILE.write_text(json.dumps(state, indent=2), encoding="utf-8")
+    try:
+        create_repo(STATE_REPO, repo_type="dataset", exist_ok=True, token=HF_MODEL_TOKEN)
+        HF_API.upload_file(
+            path_or_fileobj=str(STATE_FILE),
+            path_in_repo="training_state.json",
+            repo_id=STATE_REPO,
+            repo_type="dataset",
+        )
+    except Exception as exc:
+        print(f"⚠️ Could not update training state repo: {exc}")
+
+
+training_state = load_training_state()
+
+def mark_phase_complete(phase_number: int):
+    training_state[f"phase_{phase_number}"] = {
+        "status": "complete",
+        "timestamp": time.time(),
+    }
+    training_state["last_completed_phase"] = phase_number
+    save_training_state(training_state)
+
+
+def mark_phase_failed(phase_number: int, message: str):
+    training_state[f"phase_{phase_number}"] = {
+        "status": "failed",
+        "timestamp": time.time(),
+        "message": message,
+    }
+    save_training_state(training_state)
+
+
 def _dir_has_arrow_files(path: Path) -> bool:
     return path.exists() and any(path.glob("*.arrow"))
 
@@ -220,7 +284,10 @@ def start_training(start_option: str):
 
     option = start_option or "Auto (skip completed phases)"
     skip_completed = option.startswith("Auto")
-    start_from = 1
+    if option.startswith("Auto"):
+        start_from = max(1, training_state.get("last_completed_phase", 0) + 1)
+    else:
+        start_from = 1
     if option.startswith("Start from Phase"):
         try:
             start_from = int(option.split()[3])
@@ -279,6 +346,8 @@ def start_training(start_option: str):
                         )
                         log_f.write(skip_msg)
                         log_f.flush()
+                        if training_state.get(f"phase_{phase_number}", {}).get("status") != "complete":
+                            mark_phase_complete(phase_number)
                         continue
 
                     if phase_number < start_from and not phase_complete:
@@ -295,13 +364,29 @@ def start_training(start_option: str):
                         log_f.flush()
                         training_status["phase"] = f"PHASE {phase_number}: {phase_label}"
                         training_status["progress"] = "Already complete—skipping."
+                        if training_state.get(f"phase_{phase_number}", {}).get("status") != "complete":
+                            mark_phase_complete(phase_number)
                         continue
 
                     if not script_path.exists():
-                        training_status["progress"] = f"❌ Missing script: {script_name}"
+                        message = f"Missing script: {script_name}"
+                        training_status["progress"] = f"❌ {message}"
+                        mark_phase_failed(phase_number, message)
                         success = False
                         break
 
+                    if phase_number == 5 and not (_phase_completed(3) and _phase_completed(4)):
+                        msg = (
+                            "⚠️ Skipping evaluation: forward and retro models are not yet available on the Hub."
+                            " Complete Phases 3 and 4 before running evaluation.\n"
+                        )
+                        log_f.write(msg)
+                        log_f.flush()
+                        training_status["phase"] = f"PHASE {phase_number}: {phase_label}"
+                        training_status["progress"] = "Skipped evaluation—models missing."
+                        mark_phase_failed(phase_number, "Models missing for evaluation")
+                        continue
+
                     phase_header = f"--- Phase {phase_number}: {phase_label} ---\n"
                     log_f.write(phase_header)
                     log_f.flush()
@@ -313,24 +398,16 @@ def start_training(start_option: str):
                     )
 
                     if return_code != 0:
-                        training_status["progress"] = (
-                            f"❌ {phase_label} failed (exit code {return_code}). Check the logs above."
+                        message = (
+                            f"{phase_label} failed (exit code {return_code}). Check the logs above."
                         )
+                        training_status["progress"] = f"❌ {message}"
+                        mark_phase_failed(phase_number, message)
                         success = False
                         break
 
                     training_status["progress"] = f"✅ {phase_label} completed."
-
-                    if phase_number == 5 and not (_phase_completed(3) and _phase_completed(4)):
-                        msg = (
-                            "⚠️ Skipping evaluation: forward and retro models are not yet available on the Hub."
-                            " Complete Phases 3 and 4 before running evaluation.\n"
-                        )
-                        log_f.write(msg)
-                        log_f.flush()
-                        training_status["phase"] = f"PHASE {phase_number}: {phase_label}"
-                        training_status["progress"] = "Skipped evaluation—models missing."
-                        continue
+                    mark_phase_complete(phase_number)
 
         except Exception as exc:  # pragma: no cover - defensive logging
             success = False

src/config.py

@@ -13,6 +13,8 @@ MODELS_DIR = PROJECT_ROOT / "models"
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 MODELS_DIR.mkdir(parents=True, exist_ok=True)
 
+STATE_REPO = f"{HF_USERNAME}/ord-training-state"
+
 # Hugging Face Model and Dataset Names
 TOKENIZER_NAME = f"{HF_USERNAME}/ord-tokenizer"
 FORWARD_MODEL_NAME = f"{HF_USERNAME}/ord-forward-t5"