Update app.py

app.py

@@ -1,17 +1,18 @@
-import os, subprocess, json, pathlib, time
+import os, subprocess, json, pathlib, time, shutil
 import gradio as gr
 
-# ---------- CONSTANTS ----------
-RUN_ROOT = "/home/user/app/runs"         # where all runs are stored (visible in App Files)
-LAST_PTR = pathlib.Path(RUN_ROOT) / "LAST"  # file that stores path to the most recent run
+# ---------- CONSTANTS (visible in App Files) ----------
+RUN_ROOT = "/home/user/app/runs"          # where all runs live
+LOG_ROOT = "/home/user/app/logs"          # global logs (so we don't pre-create run dirs)
+LAST_PTR = pathlib.Path(RUN_ROOT) / "LAST" # remembers most recent run path
 os.makedirs(RUN_ROOT, exist_ok=True)
+os.makedirs(LOG_ROOT, exist_ok=True)
 
-# env helpers (with correct names)
-DEFAULT_REPO_ID = os.environ.get("REPO_ID", "")
+# ---------- ENV / HUB ----------
+DEFAULT_REPO_ID = os.environ.get("REPO_ID", "")       # e.g. "zino36/lerobot-pusht-colab"
 PUSH_DEFAULT    = os.environ.get("PUSH_TO_HUB", "true").lower() in {"1","true","yes"}
 HF_TOKEN        = os.environ.get("HF_TOKEN")
 
-# Optional: login with a Space secret named HF_TOKEN
 if HF_TOKEN:
     try:
         from huggingface_hub import login
@@ -43,21 +44,33 @@ def tail_file(path: str, n=200):
 
 # ---------- RUN DIR HELPERS ----------
 def new_run_dir():
-    d = pathlib.Path(RUN_ROOT) / f"pusht_{int(time.time())}"
-    d.mkdir(parents=True, exist_ok=True)
+    """Return a unique run dir path WITHOUT creating it (so LeRobot can create it)."""
+    base = pathlib.Path(RUN_ROOT) / f"pusht_{int(time.time())}"
+    d = base
+    i = 1
+    while d.exists():
+        d = pathlib.Path(f"{base}_{i}")
+        i += 1
     LAST_PTR.write_text(str(d))
     return str(d)
 
 def current_run_dir(user_override: str | None):
+    """Prefer user text if given, else use the LAST pointer if present."""
     if user_override and user_override.strip():
         return user_override.strip()
     if LAST_PTR.exists():
         return LAST_PTR.read_text().strip()
-    return ""  # none yet
+    return ""
 
 def has_checkpoint(run_dir: str):
+    """We consider a checkpoint present once checkpoints/last/ exists (first save is at step 500)."""
     return os.path.isdir(os.path.join(run_dir, "checkpoints", "last"))
 
+def train_log_path_for_new(run_dir: str):
+    """Write fresh-run logs to global LOG_ROOT so we don't pre-create run_dir."""
+    name = pathlib.Path(run_dir).name
+    return os.path.join(LOG_ROOT, f"{name}.train.log")
+
 def train_log_path(run_dir: str):
     return os.path.join(run_dir, "logs", "train.log")
 
@@ -67,7 +80,7 @@ def eval_log_path(run_dir: str):
 # ---------- ACTIONS ----------
 def start_training(steps, batch_size, push_to_hub, repo_id):
     run_dir = new_run_dir()
-    log = train_log_path(run_dir)
+    log = train_log_path_for_new(run_dir)
 
     push_flags = (f"--policy.push_to_hub=true --policy.repo_id='{repo_id.strip()}'"
                   if push_to_hub and repo_id.strip() else
@@ -96,7 +109,7 @@ def resume_training(extra_steps, push_to_hub, repo_id, run_dir_text):
     log = train_log_path(run_dir)
 
     if not has_checkpoint(run_dir):
-        return f"No checkpoint in {run_dir}/checkpoints/last/ yet — let the run save once (>= first 500 steps).", run_dir, tail_file(log)
+        return f"No checkpoint in {run_dir}/checkpoints/last/ yet — run at least 500 steps once.", run_dir, tail_file(log)
 
     push_flags = (f"--policy.push_to_hub=true --policy.repo_id='{repo_id.strip()}'"
                   if push_to_hub and repo_id.strip() else
@@ -148,17 +161,48 @@ def eval_latest(run_dir_text):
     msg = f"Evaluated run at: {run_dir}\nEval exited rc={rc}\n\n=== eval.log tail ===\n{tail}"
     return msg, run_dir, tail_file(elog), metrics_txt
 
+# ---------- Maintenance (list / delete runs) ----------
 def list_runs():
     root = pathlib.Path(RUN_ROOT)
     if not root.exists():
         return "(no runs)"
     rows = []
     for d in sorted(root.glob("pusht_*")):
-        size = subprocess.check_output(["bash","-lc", f"du -sh {d} | cut -f1"], text=True).strip()
+        try:
+            size = subprocess.check_output(
+                ["bash","-lc", f"du -sh {d} | cut -f1"], text=True
+            ).strip()
+        except Exception:
+            size = "?"
         ck = "✓" if has_checkpoint(str(d)) else "—"
         rows.append(f"{d.name}\t{size}\tcheckpoint:{ck}")
     return "name\tsize\tcheckpoint\n" + "\n".join(rows) if rows else "(no runs)"
 
+def delete_run_by_name(name: str):
+    name = os.path.basename((name or "").strip())
+    if not name:
+        return "Type a folder like 'pusht_1234567890'.", list_runs()
+    target = os.path.join(RUN_ROOT, name)
+    if not target.startswith(RUN_ROOT + "/"):
+        return "Refusing to delete outside runs/.", list_runs()
+    if not os.path.isdir(target):
+        return f"Folder not found: {target}", list_runs()
+    shutil.rmtree(target, ignore_errors=True)
+    # clear LAST if it pointed here
+    if LAST_PTR.exists() and LAST_PTR.read_text().strip() == target:
+        LAST_PTR.unlink(missing_ok=True)
+    return f"Deleted {target}", list_runs()
+
+def delete_all_runs():
+    if not os.path.isdir(RUN_ROOT):
+        return "(runs/ missing)", list_runs()
+    for n in os.listdir(RUN_ROOT):
+        p = os.path.join(RUN_ROOT, n)
+        if os.path.isdir(p) and n.startswith("pusht_"):
+            shutil.rmtree(p, ignore_errors=True)
+    LAST_PTR.unlink(missing_ok=True)
+    return "Deleted all pusht_* runs.", list_runs()
+
 # ---------- UI ----------
 with gr.Blocks(title="LeRobot PushT Trainer (Space)") as demo:
     gr.Markdown("# 🤖 LeRobot PushT Trainer\nTrain / Resume / Evaluate. Files persist under `/home/user/app/runs/` (see App Files).")
@@ -194,10 +238,18 @@ with gr.Blocks(title="LeRobot PushT Trainer (Space)") as demo:
     list_btn = gr.Button("📂 List runs folder")
     list_out = gr.Textbox(label="runs/ listing", lines=12)
 
+    gr.Markdown("### Maintenance")
+    del_name = gr.Textbox(label="Run folder name to delete (e.g., pusht_1699999999)")
+    del_one_btn = gr.Button("🗑️ Delete this run")
+    del_all_btn = gr.Button("🧹 Delete ALL pusht_* runs")
+
+    # Wiring
     start_btn.click(start_training, inputs=[steps, batch, push_to_hub, repo_id], outputs=[start_out, run_dir_view, train_log])
     resume_btn.click(resume_training, inputs=[extra_steps, push_to_hub, repo_id, run_dir_text], outputs=[resume_out, run_dir_view, resume_log])
     eval_btn.click(eval_latest, inputs=[run_dir_text], outputs=[eval_out, run_dir_view, eval_log, metrics_box])
     list_btn.click(list_runs, outputs=list_out)
+    del_one_btn.click(delete_run_by_name, inputs=del_name, outputs=[list_out, list_out])
+    del_all_btn.click(delete_all_runs, outputs=[list_out, list_out])
 
 if __name__ == "__main__":
     demo.launch()