Added logic to download csv files first then generate Master Excel file

app.py

@@ -670,7 +670,8 @@ def write_master_workbook(path, sheets_in_order, row_limit=1_048_575):
 # MAIN PIPELINE
 # =====================================================================
 def run_pipeline(uploaded_files, reference_file, blank_cities_not_in_list,
-                 filter_dominant_county, cluster_threshold, replace_threshold):
+                 filter_dominant_county, cluster_threshold, replace_threshold,
+                 build_master=True):
     t_start = time.time()
     log = []
     if not uploaded_files:
@@ -977,33 +978,13 @@ def run_pipeline(uploaded_files, reference_file, blank_cities_not_in_list,
     summary_df.to_csv(summary_path, index=False)
     outputs.append(str(summary_path))
 
-    # ---- Master bundle: Excel workbook OR zip of CSVs (see MASTER_OUTPUT) -
-    _step(log, f"Building master ({MASTER_OUTPUT})...")
-    if MASTER_OUTPUT == "csv":
-        # Bundle the already-written CSVs into one zip. No Excel writing -> fast.
-        master_path = out_dir / "Master_Dataset.zip"
-        member_names = ["FIPS_Validation_Audit_ChangedRecords.csv", "cluster_df.csv",
-                        "final_cities.csv", "final_zips.csv", "final_census_tracts.csv",
-                        "Summary.csv"]
-        members = [out_dir / n for n in member_names if (out_dir / n).exists()]
-        members += sorted(out_dir.glob("OutputFile_*.csv"))   # full output (chunked)
-        with zipfile.ZipFile(master_path, "w", zipfile.ZIP_DEFLATED) as zf:
-            for p in members:
-                zf.write(p, arcname=p.name)
-    else:
-        # Every result on its own sheet; the full output splits across
-        # Output_70_95_RapidFuzz_1, _2, ... if it exceeds Excel's row cap.
-        master_path = out_dir / "Master_Dataset.xlsx"
-        write_master_workbook(master_path, [
-            ("Output_70_95_RapidFuzz", output_df),
-            ("ValidationAuditResults", audit),
-            ("cluster_df", cluster_df),
-            ("final_cities", final_cities),
-            ("final_zips", final_zips),
-            ("final_census_tracts", final_ct),
-            ("Summary", summary_df),
-        ])
-    outputs.append(str(master_path))
+    # ---- Master (slow on huge states) -- optional / best-effort ------
+    # Built from the CSVs already on disk (same as streaming), so it can be
+    # skipped here and run separately after the individual files are secured.
+    if build_master:
+        _step(log, f"Building master ({MASTER_OUTPUT})...")
+        outputs.append(build_master_bundle(out_dir, MASTER_OUTPUT))
+        _step(log, "Master built")
 
     elapsed_min = (time.time() - t_start) / 60
     log.append(f"Done. {len(outputs)} files ready.")
@@ -1069,20 +1050,41 @@ def run_from_dataset(input_repo, output_repo):
         except Exception as e:
             return f"Failed to download '{fn}' from '{input_repo}': {e}", []
 
-    log, outputs = process(local, None)
+    # Process WITHOUT the master first, so the individual files exist quickly.
+    log, outputs = process(local, None, build_master=False)
 
+    stamp = time.strftime("%Y%m%d_%H%M%S")
+    dest = f"https://huggingface.co/datasets/{output_repo}/tree/main/runs/{stamp}" if output_repo else ""
+
+    def _upload(paths, api):
+        for p in paths:
+            api.upload_file(path_or_fileobj=p, path_in_repo=f"runs/{stamp}/{Path(p).name}",
+                            repo_id=output_repo, repo_type="dataset")
+
+    api = None
     if output_repo and outputs:
         try:
             api = HfApi(token=token)
-            stamp = time.strftime("%Y%m%d_%H%M%S")
-            for p in outputs:
-                api.upload_file(path_or_fileobj=p, path_in_repo=f"runs/{stamp}/{Path(p).name}",
-                                repo_id=output_repo, repo_type="dataset")
-            log += (f"\nUploaded {len(outputs)} file(s) to dataset '{output_repo}' "
-                    f"under runs/{stamp}/. Retrieve them there:\n"
-                    f"https://huggingface.co/datasets/{output_repo}/tree/main/runs/{stamp}")
+            _upload(outputs, api)                       # <-- SECURE the CSVs first
+            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}")
         except Exception as e:
-            log += f"\nWARNING: processed OK but upload to '{output_repo}' failed: {e}"
+            log += f"\nWARNING: processed OK but uploading core files failed: {e}"
+
+    # Master last / best-effort. If it dies on a huge state, the files above survive.
+    if outputs:
+        try:
+            master = build_master_bundle(Path(outputs[0]).parent, MASTER_OUTPUT)
+            outputs.append(master)
+            if api is not None:
+                _upload([master], api)
+                log += f"\nMaster built and uploaded: {Path(master).name}"
+            else:
+                log += f"\nMaster built: {Path(master).name}"
+        except Exception as e:
+            log += (f"\nMaster build/upload FAILED -- the individual files are already "
+                    f"saved in the dataset, so nothing is lost: {e}")
     return log, outputs
 
 
@@ -1322,7 +1324,8 @@ def _split_by_county(uploaded_files, out_dir, log, chunksize=STREAM_CHUNK_ROWS):
 
 
 def run_streaming(uploaded_files, reference_file, blank_cities_not_in_list,
-                  filter_dominant_county, cluster_threshold, replace_threshold):
+                  filter_dominant_county, cluster_threshold, replace_threshold,
+                  build_master=True):
     """Per-county streaming pipeline: bounded memory, identical results."""
     t_start = time.time()
     log = []
@@ -1407,24 +1410,11 @@ def run_streaming(uploaded_files, reference_file, blank_cities_not_in_list,
     outputs.append(str(summary_path))
     _step(log, f"Summary: {len(summary_df)} county row(s)")
 
-    # ---- Master ------------------------------------------------------
-    _step(log, f"Building master ({MASTER_OUTPUT})...")
-    if MASTER_OUTPUT == "csv":
-        master_path = out_dir / "Master_Dataset.zip"
-        member_names = ["FIPS_Validation_Audit_ChangedRecords.csv", "cluster_df.csv",
-                        "final_cities.csv", "final_zips.csv", "final_census_tracts.csv",
-                        "Summary.csv"]
-        members = [out_dir / n for n in member_names if (out_dir / n).exists()]
-        members += [Path(p) for p in output_parts]
-        with zipfile.ZipFile(master_path, "w", zipfile.ZIP_DEFLATED) as zf:
-            for p in members:
-                zf.write(p, arcname=Path(p).name)
-    else:
-        # Stream the big Output sheet from the combined CSV; small sheets load fully.
-        master_path = out_dir / "Master_Dataset.xlsx"
-        _write_master_streaming(master_path, output_parts, audit_path, clu_path,
-                                fc_path, fz_path, fct_path, summary_df)
-    outputs.append(str(master_path))
+    # ---- Master (slow on huge states) -- optional / best-effort ------
+    if build_master:
+        _step(log, f"Building master ({MASTER_OUTPUT})...")
+        outputs.append(build_master_bundle(out_dir, MASTER_OUTPUT))
+        _step(log, "Master built")
 
     elapsed_min = (time.time() - t_start) / 60
     log.append(f"Done. {len(outputs)} files ready.")
@@ -1499,13 +1489,50 @@ def _write_master_streaming(path, output_parts, audit_path, clu_path,
     wb.close()
 
 
-def process(uploaded_files, reference_file):
-    """Dispatcher used by the UI: streaming or in-memory per PROCESS_MODE."""
+def build_master_bundle(out_dir, master_output):
+    """Build the master (zip or xlsx) from the CSV outputs already on disk in
+    out_dir. Kept SEPARATE from the pipeline so callers can secure the fast CSVs
+    first and run this slow step last / best-effort -- if it times out on a huge
+    state, the individual files are already saved. Returns the master path."""
+    out_dir = Path(out_dir)
+    output_parts = [str(p) for p in sorted(out_dir.glob("OutputFile_*.csv"))]
+    audit_path = out_dir / "FIPS_Validation_Audit_ChangedRecords.csv"
+    clu_path = out_dir / "cluster_df.csv"
+    fc_path = out_dir / "final_cities.csv"
+    fz_path = out_dir / "final_zips.csv"
+    fct_path = out_dir / "final_census_tracts.csv"
+    summary_path = out_dir / "Summary.csv"
+
+    if master_output == "csv":
+        master_path = out_dir / "Master_Dataset.zip"
+        member_names = ["FIPS_Validation_Audit_ChangedRecords.csv", "cluster_df.csv",
+                        "final_cities.csv", "final_zips.csv", "final_census_tracts.csv",
+                        "Summary.csv"]
+        members = [out_dir / n for n in member_names if (out_dir / n).exists()]
+        members += [Path(p) for p in output_parts]
+        with zipfile.ZipFile(master_path, "w", zipfile.ZIP_DEFLATED) as zf:
+            for p in members:
+                zf.write(p, arcname=Path(p).name)
+    else:
+        master_path = out_dir / "Master_Dataset.xlsx"
+        summary_df = (pd.read_csv(summary_path, low_memory=False)
+                      if summary_path.exists() else pd.DataFrame())
+        _write_master_streaming(master_path, output_parts, audit_path, clu_path,
+                                fc_path, fz_path, fct_path, summary_df)
+    return str(master_path)
+
+
+def process(uploaded_files, reference_file, build_master=True):
+    """Dispatcher used by the UI: streaming or in-memory per PROCESS_MODE.
+    build_master=False returns just the individual files (no master), so a caller
+    can save/upload those first and build the slow master separately."""
     if PROCESS_MODE == "stream":
         return run_streaming(uploaded_files, reference_file, STRICT_BLANK_CITIES,
-                             FILTER_DOMINANT_COUNTY, CLUSTER_THRESHOLD, REPLACE_THRESHOLD)
+                             FILTER_DOMINANT_COUNTY, CLUSTER_THRESHOLD, REPLACE_THRESHOLD,
+                             build_master=build_master)
     return run_pipeline(uploaded_files, reference_file, STRICT_BLANK_CITIES,
-                        FILTER_DOMINANT_COUNTY, CLUSTER_THRESHOLD, REPLACE_THRESHOLD)
+                        FILTER_DOMINANT_COUNTY, CLUSTER_THRESHOLD, REPLACE_THRESHOLD,
+                        build_master=build_master)
 
 
 # =====================================================================
@@ -1577,8 +1604,20 @@ with gr.Blocks() as demo:
             files_out = gr.File(label="Output files", file_count="multiple")
 
     def _run(files, reference):
-        # All settings are fixed (see constants at top). No user-facing options.
-        return process(files, reference)
+        # Produce the individual files first and show them for download, THEN build
+        # the (slow-on-huge-files) master. If the master step is interrupted, the
+        # files below are already downloadable.
+        log, outputs = process(files, reference, build_master=False)
+        yield log + "\n\nFiles below are ready to download. Building master...", outputs
+        if outputs:
+            try:
+                master = build_master_bundle(Path(outputs[0]).parent, MASTER_OUTPUT)
+                outputs = outputs + [master]
+                yield log + f"\n\nMaster ready: {Path(master).name}", outputs
+            except Exception as e:
+                yield log + f"\n\nFiles are ready. Master build failed: {e}", outputs
+        else:
+            yield log, outputs
 
     run_btn.click(
         _run,