Upload app.py

app.py

@@ -74,6 +74,16 @@ STREAM_CHUNK_ROWS = 400_000   # rows per chunk while splitting the input by coun
 # on big states, so it's OFF by default. Set True if you actually need that file.
 WRITE_RAWFILE = False
 
+# ---- Dataset auto-cleanup (Run from dataset) -------------------------------
+# HF datasets are git repos: deleting a file leaves its bytes in git HISTORY, so
+# storage is only reclaimed by SQUASHING the history. These options keep the two
+# datasets small automatically so you never hit the private-storage quota again.
+# All of this runs ONLY after results have been safely uploaded, so nothing is
+# lost if a run fails.
+DELETE_INPUT_AFTER_RUN = True   # delete the processed input file(s) from the input dataset
+KEEP_LATEST_OUTPUT_RUNS = 1     # in the output dataset, keep only the N newest runs/ folders (0 = keep all)
+SQUASH_HISTORY = True           # after deleting, squash both datasets' history to reclaim the storage
+
 # Threshold for the RURAL/URBAN label (records <= this -> RURAL, else URBAN).
 RURAL_MAX_RECORDS = 10000
 
@@ -1154,6 +1164,69 @@ def list_dataset_inputs(repo_id, token=None):
     return sorted(f for f in files if f.lower().endswith((".csv", ".xlsx", ".xls")))
 
 
+def _cleanup_datasets(api, input_repo, output_repo, processed_inputs, current_stamp):
+    """Keep the datasets small so the private-storage quota is never hit. Runs
+    ONLY after results are safely uploaded. Returns a short log string.
+    Steps (each guarded by its config flag):
+      1. delete the input file(s) we just processed from the input dataset
+      2. keep only the newest KEEP_LATEST_OUTPUT_RUNS runs/ folders in the output dataset
+      3. squash both datasets' git history -- REQUIRED to actually reclaim bytes,
+         since a plain delete leaves the old versions in history."""
+    from huggingface_hub import list_repo_files
+    notes = []
+    squashed_any = False
+
+    # 1) delete processed inputs
+    if DELETE_INPUT_AFTER_RUN and input_repo and processed_inputs:
+        deleted = 0
+        for fn in processed_inputs:
+            try:
+                api.delete_file(path_in_repo=fn, repo_id=input_repo,
+                                repo_type="dataset",
+                                commit_message=f"Auto-cleanup: remove processed input {fn}")
+                deleted += 1
+            except Exception as e:
+                notes.append(f"  (could not delete input '{fn}': {e})")
+        if deleted:
+            notes.append(f"  Deleted {deleted} processed input file(s) from '{input_repo}'.")
+
+    # 2) prune old run folders in the output dataset -- keep the newest
+    #    KEEP_LATEST_OUTPUT_RUNS folders TOTAL, and the run we just wrote always
+    #    counts as one of them (so it is never deleted).
+    if KEEP_LATEST_OUTPUT_RUNS and output_repo:
+        try:
+            files = list_repo_files(output_repo, repo_type="dataset")
+            stamps = sorted({p.split("/")[1] for p in files
+                             if p.startswith("runs/") and len(p.split("/")) > 2})
+            others = [s for s in stamps if s != current_stamp]   # oldest..newest
+            keep_others = others[-(KEEP_LATEST_OUTPUT_RUNS - 1):] if KEEP_LATEST_OUTPUT_RUNS > 1 else []
+            old = [s for s in others if s not in keep_others]
+            for s in old:
+                try:
+                    api.delete_folder(path_in_repo=f"runs/{s}", repo_id=output_repo,
+                                      repo_type="dataset",
+                                      commit_message=f"Auto-cleanup: remove old run {s}")
+                except Exception as e:
+                    notes.append(f"  (could not delete old run '{s}': {e})")
+            if old:
+                notes.append(f"  Pruned {len(old)} older run folder(s) from '{output_repo}'.")
+        except Exception as e:
+            notes.append(f"  (could not list runs for pruning: {e})")
+
+    # 3) squash history to actually reclaim the storage
+    if SQUASH_HISTORY:
+        for repo in filter(None, {input_repo, output_repo}):
+            try:
+                api.super_squash_history(repo_id=repo, repo_type="dataset")
+                squashed_any = True
+            except Exception as e:
+                notes.append(f"  (could not squash history for '{repo}': {e})")
+        if squashed_any:
+            notes.append("  Squashed dataset history -> storage reclaimed.")
+
+    return ("\n\nAuto-cleanup:\n" + "\n".join(notes)) if notes else ""
+
+
 def run_from_dataset(input_repo, output_repo):
     """Download every supported file from input_repo, run the pipeline, and (if
     output_repo is set) upload the results back to that dataset. Returns
@@ -1194,10 +1267,12 @@ def run_from_dataset(input_repo, output_repo):
                             repo_id=output_repo, repo_type="dataset")
 
     api = None
+    core_uploaded = False
     if output_repo and outputs:
         try:
             api = HfApi(token=token)
             _upload(outputs, api)                       # <-- SECURE the CSVs first
+            core_uploaded = True
             log += (f"\nUploaded {len(outputs)} core file(s) to '{output_repo}' under "
                     f"runs/{stamp}/ BEFORE building the master. If the master step "
                     f"fails, these are already saved.\n{dest}")
@@ -1218,6 +1293,13 @@ def run_from_dataset(input_repo, output_repo):
             log += (f"\nMaster build/upload FAILED -- the individual files are already "
                     f"saved in the dataset, so nothing is lost: {e}")
 
+    # Auto-cleanup runs ONLY once results are safely uploaded, so nothing is lost.
+    if api is not None and core_uploaded:
+        try:
+            log += _cleanup_datasets(api, input_repo, output_repo, filenames, stamp)
+        except Exception as e:
+            log += f"\nAuto-cleanup skipped (non-fatal): {e}"
+
     # A clean, clickable list of every file in the dataset, so nobody has to hunt.
     if output_repo and outputs:
         base = f"https://huggingface.co/datasets/{output_repo}/blob/main/runs/{stamp}"