142 lines
3.9 KiB
Python
142 lines
3.9 KiB
Python
from pathlib import Path
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import threading
|
|
import wandb
|
|
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()
|
|
|
|
print_lock = threading.Lock()
|
|
|
|
|
|
def safe_print(*args, **kwargs):
|
|
with print_lock:
|
|
print(*args, **kwargs)
|
|
|
|
|
|
def download_file(file, run_dir):
|
|
target = run_dir / file.name
|
|
|
|
try:
|
|
# skip existing files
|
|
if target.exists():
|
|
return f"SKIP FILE {target}"
|
|
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
file.download(root=run_dir, replace=False)
|
|
|
|
return f"DONE FILE {target}"
|
|
|
|
except Exception as 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):
|
|
run_dir = root / f"{run.name}"
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
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:
|
|
futures = [executor.submit(download_file, file, run_dir) for file in files]
|
|
|
|
for future in as_completed(futures):
|
|
safe_print(future.result())
|
|
|
|
# -------------------------
|
|
# 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}")
|
|
|
|
|
|
def main(entity: str, project: str, root: Path):
|
|
root.mkdir(exist_ok=True)
|
|
|
|
runs = list(api.runs(f"{entity}/{project}"))
|
|
|
|
safe_print(f"Found {len(runs)} runs")
|
|
|
|
with ThreadPoolExecutor(max_workers=MAX_RUN_WORKERS) as executor:
|
|
futures = [executor.submit(download_run, run, root) for run in runs]
|
|
|
|
for future in as_completed(futures):
|
|
try:
|
|
future.result()
|
|
except Exception as e:
|
|
safe_print("RUN FAILED:", e)
|
|
|
|
safe_print("\nAll downloads complete.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--entity", type=str, default="SEL3-2026-Groep-4")
|
|
parser.add_argument("--project", type=str, required=True)
|
|
parser.add_argument("--root", type=str, default="runs")
|
|
args = parser.parse_args()
|
|
|
|
root = Path(args.root)
|
|
main(entity=args.entity, project=args.project, root=root)
|