From 03b1b3397490038281f3aa04afd92ca64a7eefdd Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Thu, 14 May 2026 21:42:50 +0200 Subject: [PATCH] feat(downloader): added downloading of artifacts + moved script to tool directory --- scripts/plots/analyze_convergence.py | 30 +++++--- scripts/{ => tools}/download_wandb_project.py | 69 +++++++++++++++++-- 2 files changed, 84 insertions(+), 15 deletions(-) rename scripts/{ => tools}/download_wandb_project.py (51%) diff --git a/scripts/plots/analyze_convergence.py b/scripts/plots/analyze_convergence.py index 2612bb9..08f905d 100644 --- a/scripts/plots/analyze_convergence.py +++ b/scripts/plots/analyze_convergence.py @@ -45,19 +45,22 @@ class Columns(str, Enum): """Column names expected in every evaluation CSV.""" ARCH = "architecture" - TIMESTEPS = "total_trained_timesteps" - REWARD = "accumulated_reward" + TIMESTEPS = "trained_timesteps" + REWARD = "eval_return" VELOCITY = "velocity" + EVAL_STEPS = "eval_steps" + FINAL_XY_DIST = "final_xy_dist" + INITIAL_XY_DIST = "initial_xy_dist" + REACHED_TARGET = "reached_target" # Maps architecture display names to the path of their evaluation CSV. # Update these paths once real evaluation data is available. FILE_MAPPING: dict[str, str] = { - "centralized 2 arms": "runs/dummy/dummy_centralized_2_arms.csv", - "centralized 5 arms": "runs/dummy/dummy_centralized_5_arms.csv", - "decentralized fully connected": "runs/dummy/dummy_decentralized_fully_connected.csv", - "decentralized ring-level": "runs/dummy/dummy_decentralized_ring-level.csv", - "decentralized segment-level": "runs/dummy/dummy_decentralized_segment-level.csv", + # "centralized 2 arms": "runs/dummy/dummy_centralized_2_arms.csv", + "centralized 5 arms": "runs/final-v2-centralized/checkpoint_evaluation.csv", + "decentralized fully connected": "runs/final-v2-fully-conn/checkpoint_evaluation.csv", + "decentralized ring-level": "runs/final-v2-ring/checkpoint_evaluation.csv", } # Architecture profiles for dummy data generation: (max_reward, max_velocity, sigmoid_speed) @@ -108,7 +111,13 @@ def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame: Loads one CSV per architecture, injects the architecture name as a column, and returns the combined DataFrame with only the required columns. """ - required = [Columns.TIMESTEPS, Columns.REWARD, Columns.VELOCITY] + required = [ + Columns.TIMESTEPS, + Columns.REWARD, + Columns.INITIAL_XY_DIST, + Columns.FINAL_XY_DIST, + Columns.EVAL_STEPS, + ] dfs = [] for arch_name, filepath in file_mapping.items(): @@ -125,6 +134,10 @@ def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame: df = df[required].copy() df[Columns.ARCH] = arch_name + df[Columns.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[ + Columns.EVAL_STEPS + ] + dfs.append(df) return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame() @@ -303,6 +316,7 @@ def plot_results(df: pd.DataFrame, results: pd.DataFrame, output_dir: str, **kwa def obtain_data() -> pd.DataFrame: """Resolves the file mapping, falling back to generated dummy CSVs if needed.""" global USING_DUMMY_DATA + if not any(os.path.exists(p) for p in FILE_MAPPING.values()): logger.info("No real evaluation files found. Generating dummy CSVs at expected locations.") generate_dummy_csvs(FILE_MAPPING) diff --git a/scripts/download_wandb_project.py b/scripts/tools/download_wandb_project.py similarity index 51% rename from scripts/download_wandb_project.py rename to scripts/tools/download_wandb_project.py index 7c04ad1..bf2fb11 100644 --- a/scripts/download_wandb_project.py +++ b/scripts/tools/download_wandb_project.py @@ -7,6 +7,7 @@ import argparse # tune these depending on network / W&B limits MAX_RUN_WORKERS = 8 MAX_FILE_WORKERS = 16 +MAX_ARTIFACT_WORKERS = 8 api = wandb.Api() @@ -20,19 +21,42 @@ def safe_print(*args, **kwargs): def download_file(file, run_dir): target = run_dir / file.name + try: # skip existing files if target.exists(): - return f"SKIP {target}" + return f"SKIP FILE {target}" target.parent.mkdir(parents=True, exist_ok=True) file.download(root=run_dir, replace=False) - return f"DONE {target}" + return f"DONE FILE {target}" except Exception as e: - return f"FAIL {target}: {e}" + return f"FAIL FILE {target}: {e}" + + +def sanitize_artifact_name(name: str): + return name.replace(":", "_") + + +def download_artifact(artifact, artifact_root): + try: + artifact_name = sanitize_artifact_name(artifact.name) + artifact_dir = artifact_root / artifact_name + + if artifact_dir.exists() and any(artifact_dir.iterdir()): + return f"SKIP ARTIFACT {artifact.name}" + + artifact_dir.mkdir(parents=True, exist_ok=True) + + artifact.download(root=artifact_dir) + + return f"DONE ARTIFACT {artifact.name}" + + except Exception as e: + return f"FAIL ARTIFACT {artifact.name}: {e}" def download_run(run, root): @@ -41,6 +65,9 @@ def download_run(run, root): safe_print(f"\n=== {run.name} ({run.id}) ===") + # ------------------------- + # Download regular run files + # ------------------------- files = list(run.files()) with ThreadPoolExecutor(max_workers=MAX_FILE_WORKERS) as executor: @@ -49,10 +76,38 @@ def download_run(run, root): for future in as_completed(futures): safe_print(future.result()) - # OPTIONAL: download artifacts too - # for artifact in run.logged_artifacts(): - # artifact_dir = run_dir / "artifacts" / artifact.name - # artifact.download(root=artifact_dir) + # ------------------------- + # Download logged artifacts + # ------------------------- + artifact_root = run_dir / "artifacts" + + try: + artifacts = list(run.logged_artifacts()) + safe_print(f"Found {len(artifacts)} artifacts for {run.name}") + + with ThreadPoolExecutor(max_workers=MAX_ARTIFACT_WORKERS) as executor: + futures = [ + executor.submit(download_artifact, artifact, artifact_root) + for artifact in artifacts + ] + + for future in as_completed(futures): + safe_print(future.result()) + + except Exception as e: + safe_print(f"Artifact download failed for {run.name}: {e}") + + # ------------------------- + # OPTIONAL: download used/input artifacts + # ------------------------- + # try: + # used_artifacts = list(run.used_artifacts()) + # used_root = run_dir / "used_artifacts" + # + # for artifact in used_artifacts: + # download_artifact(artifact, used_root) + # except Exception as e: + # safe_print(f"Used artifact download failed: {e}") safe_print(f"Finished {run.name}")