#!/usr/bin/env python3 """Download pre-computed model artifacts from GitHub releases. These artifacts contain registry databases, model predictions, and backtest results that allow readers to run strategy notebooks (Ch16-20) and insight notebooks (Ch11-15) without first training all models. The artifacts are published as a GitHub release asset and are added as the case-study chapters roll out; until that release lands, the notebooks still run end to end from scratch (the artifacts only skip retraining). Usage: python scripts/download_artifacts.py # all case studies python scripts/download_artifacts.py --cs etfs # single case study python scripts/download_artifacts.py --cs etfs --force # re-download python scripts/download_artifacts.py --list # show available """ import argparse import os import shutil import subprocess import sys import tarfile import urllib.request from pathlib import Path REPO_ROOT = Path(__file__).resolve().parent.parent # GitHub release configuration GITHUB_REPO = "stefan-jansen/machine-learning-for-trading" RELEASE_TAG = "v3.0.0-artifacts" BASE_URL = f"https://github.com/{GITHUB_REPO}/releases/download/{RELEASE_TAG}" CASE_STUDIES = [ "etfs", "crypto_perps_funding", "nasdaq100_microstructure", "sp500_equity_option_analytics", "us_firm_characteristics", "fx_pairs", "cme_futures", "sp500_options", "us_equities_panel", ] RELEASE_HINT = ( f"The pre-computed artifacts release ({RELEASE_TAG}) may not be published yet.\n" "The artifacts are added as the case-study chapters roll out; until then every\n" "notebook still runs end to end from scratch — the artifacts only skip retraining.\n" f"Check the latest releases at https://github.com/{GITHUB_REPO}/releases" ) def _get_github_token() -> str | None: """Get GitHub token from env or gh CLI.""" token = os.environ.get("GITHUB_TOKEN") if token: return token try: result = subprocess.run(["gh", "auth", "token"], capture_output=True, text=True, check=True) return result.stdout.strip() except (subprocess.CalledProcessError, FileNotFoundError): return None def _download_with_gh(asset_name: str, dest: Path, desc: str) -> bool: """Download using gh CLI (handles auth automatically).""" dest.parent.mkdir(parents=True, exist_ok=True) try: subprocess.run( [ "gh", "release", "download", RELEASE_TAG, "--repo", GITHUB_REPO, "--pattern", asset_name, "--dir", str(dest.parent), "--clobber", ], check=True, capture_output=True, text=True, ) downloaded = dest.parent / asset_name if downloaded != dest: downloaded.rename(dest) print(f" {desc}: done ({dest.stat().st_size / 1024 / 1024:.0f} MB)") return True except (subprocess.CalledProcessError, FileNotFoundError) as e: print(f" {desc}: gh download failed ({e})") return False def download_file(url: str, dest: Path, desc: str) -> bool: """Download a file with progress reporting. Uses token auth if available.""" token = _get_github_token() try: req = urllib.request.Request(url) if token: req.add_header("Authorization", f"token {token}") req.add_header("Accept", "application/octet-stream") with urllib.request.urlopen(req) as response: total = int(response.headers.get("Content-Length", 0)) downloaded = 0 chunk_size = 1024 * 1024 # 1MB dest.parent.mkdir(parents=True, exist_ok=True) with open(dest, "wb") as f: while True: chunk = response.read(chunk_size) if not chunk: break f.write(chunk) downloaded += len(chunk) if total: pct = downloaded * 100 // total mb = downloaded / 1024 / 1024 total_mb = total / 1024 / 1024 print( f"\r {desc}: {mb:.0f}/{total_mb:.0f} MB ({pct}%)", end="", flush=True, ) print() return True except urllib.error.HTTPError as e: if e.code in (401, 403, 404) and shutil.which("gh"): return _download_with_gh(Path(url).name, dest, desc) if e.code in (401, 403, 404): print(f"\r {desc}: FAILED ({e.code} {e.reason}) — likely missing authentication") else: print(f"\r {desc}: FAILED ({e.code} {e.reason})") return False except Exception as e: print(f"\r {desc}: FAILED ({e})") return False def extract_tarball(tarball: Path, extract_to: Path) -> int: """Extract tarball and return number of files extracted.""" count = 0 with tarfile.open(tarball, "r:gz") as tar: tar.extractall(path=extract_to, filter="data") count = len(tar.getmembers()) return count def has_artifacts(cs_id: str) -> bool: """Check if a case study already has artifacts.""" cs_dir = REPO_ROOT / "case_studies" / cs_id / "run_log" return (cs_dir / "registry.db").exists() def download_case_study(cs_id: str, force: bool = False) -> bool: """Download and extract artifacts for one case study.""" if has_artifacts(cs_id) and not force: print(f" {cs_id}: already has artifacts (use --force to re-download)") return True tarball_name = f"{cs_id}.tar.gz" url = f"{BASE_URL}/{tarball_name}" tmp_path = REPO_ROOT / ".cache" / tarball_name if not download_file(url, tmp_path, cs_id): return False print(" Extracting...", end=" ", flush=True) n = extract_tarball(tmp_path, REPO_ROOT) print(f"{n} files") # Clean up tarball tmp_path.unlink() return True def main(): parser = argparse.ArgumentParser( description="Download pre-computed model artifacts from GitHub releases" ) parser.add_argument("--cs", "--case-study", help="Single case study ID") parser.add_argument("--force", action="store_true", help="Re-download even if artifacts exist") parser.add_argument("--list", action="store_true", help="List available case studies") args = parser.parse_args() if args.list: print("Available case studies:") for cs in CASE_STUDIES: status = "installed" if has_artifacts(cs) else "not installed" print(f" {cs:40s} [{status}]") return cs_list = [args.cs] if args.cs else CASE_STUDIES # Validate for cs in cs_list: if cs not in CASE_STUDIES: print(f"Unknown case study: {cs}") print(f"Available: {', '.join(CASE_STUDIES)}") sys.exit(1) print(f"Downloading artifacts for {len(cs_list)} case study(ies)") print(f"Source: {BASE_URL}\n") # Ensure cache dir (REPO_ROOT / ".cache").mkdir(exist_ok=True) success = 0 for cs_id in cs_list: if download_case_study(cs_id, force=args.force): success += 1 # Clean up cache dir cache = REPO_ROOT / ".cache" if cache.exists() and not any(cache.iterdir()): cache.rmdir() print(f"\nDone: {success}/{len(cs_list)} case studies ready.") if success < len(cs_list): print() print(RELEASE_HINT) sys.exit(1) if __name__ == "__main__": main()