sync_epoch.py

import os
import json
import pandas as pd
import sys
import glob

# Usage: python sync_epoch.py <EPOCH_ID>

BUCKET_ROOT = "x-mas-results"
LOCAL_BASE_DIR = "./results"


def main():
    if len(sys.argv) < 2:
        print("Please provide an Epoch ID (e.g., 1)")
        sys.exit(1)

    epoch_id = sys.argv[1]

    # Remote path: gs://x-mas-results/epoch_1/
    remote_prefix = f"gs://{BUCKET_ROOT}/epoch_{epoch_id}"

    # Local path: ./results_data/epoch_1/
    local_epoch_dir = os.path.join(LOCAL_BASE_DIR, f"epoch_{epoch_id}")
    os.makedirs(local_epoch_dir, exist_ok=True)

    output_csv = os.path.join(LOCAL_BASE_DIR, f"epoch_{epoch_id}_merged.csv")

    print(f"Downloading results from {remote_prefix}...")
    # -n prevents overwriting files that exist with the EXACT name
    # BUT since we rename files locally, gsutil will re-download the original names (result_X.json)
    exit_code = os.system(f"gsutil -m cp -n {remote_prefix}/*.json {local_epoch_dir}/")

    if exit_code != 0:
        print("Warning: gsutil returned non-zero. Job might not be finished or bucket is empty.")

    print("Processing and renaming JSON files...")
    rows = []

    # Get list of all json files currently in the directory
    json_files = glob.glob(f"{local_epoch_dir}/*.json")

    for file_path in json_files:
        try:
            with open(file_path, 'r') as f:
                data = json.load(f)

            # --- 1. Extract Data ---
            name_val = data.get("name", "unknown")

            # Extract N from name (e.g., "84_7_7_0.8-0.0" -> 84)
            try:
                n_val = int(name_val.split('_')[0])
            except (ValueError, IndexError):
                print(f"Warning: Could not parse N from name '{name_val}' in file {file_path}")
                n_val = None

            # Extract Solution Metrics
            solution = data.get("solution", {})
            # If solution is None (optimization failed), use empty dict
            if solution is None:
                solution = {}

            size_val = solution.get("strip_width")  # "solution.strip_width" as size
            density_val = solution.get("density")  # "solution.density"
            start_height_val = data.get("strip_height")  # "strip_height" as start_height

            # --- 2. Rename File ---
            # New filename format: sol_{name}.json
            new_filename = f"sol_{name_val}.json"
            new_file_path = os.path.join(local_epoch_dir, new_filename)

            # Rename if the filename doesn't already match
            # (Prevents errors if running script multiple times)
            if os.path.abspath(file_path) != os.path.abspath(new_file_path):
                os.rename(file_path, new_file_path)

            # --- 3. Collect Row Data ---
            row = {
                "epoch": epoch_id,
                "N": n_val,
                "name": name_val,
                "size": size_val,
                "start_height": start_height_val,
                "density": density_val,
                "filename": f"epoch_{epoch_id}/{new_filename}"
            }
            rows.append(row)

        except Exception as e:
            print(f"Skipping broken file {file_path}: {e}")

    # --- 4. Save CSV ---
    if rows:
        df = pd.DataFrame(rows)
        df.drop_duplicates(subset=['name'], keep='last', inplace=True)
        # Enforce column order
        cols = ["epoch", "N", "name", "size", "start_height", "density", "filename"]
        # Ensure columns exist even if data was missing
        for c in cols:
            if c not in df.columns:
                df[c] = None

        df = df[cols]
        df.sort_values(by="N", inplace=True)

        df.to_csv(output_csv, index=False)
        print(f"Success! Processed {len(df)} files. Saved to {output_csv}")
    else:
        print("No valid data found to merge.")


if __name__ == "__main__":
    main()