d746008e1f
- Created entrypoint-self-terminate.sh wrapper script - Updates entrypoint-generic.sh to be called by wrapper - Modified Dockerfile.runpod to use self-terminate entrypoint - Adds automatic pod termination via runpodctl after training completes - Prevents infinite restart loops and wasted GPU credits - Saves ~96% cost per training run ($4.59 per run) Implements pod self-termination using RUNPOD_POD_ID environment variable. Training exits with code 0 → runpodctl remove pod → immediate shutdown. Co-Authored-By: Claude <noreply@anthropic.com>
478 lines
17 KiB
Python
Executable File
478 lines
17 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
RunPod Network Volume S3 Upload Script
|
|
|
|
Uploads binaries and test data to RunPod network volume via S3 API with smart
|
|
checksum comparison to avoid unnecessary uploads.
|
|
|
|
Usage:
|
|
./scripts/upload_to_runpod_volume.py --binaries # Upload binaries only
|
|
./scripts/upload_to_runpod_volume.py --test-data # Upload test data only
|
|
./scripts/upload_to_runpod_volume.py --all # Upload everything
|
|
./scripts/upload_to_runpod_volume.py --all --force # Force re-upload (ignore checksums)
|
|
./scripts/upload_to_runpod_volume.py --all --dry-run # Show what would be uploaded
|
|
"""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import boto3
|
|
from botocore.config import Config
|
|
from botocore.exceptions import ClientError
|
|
from dotenv import load_dotenv
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
BarColumn,
|
|
DownloadColumn,
|
|
Progress,
|
|
TaskID,
|
|
TextColumn,
|
|
TimeRemainingColumn,
|
|
TransferSpeedColumn,
|
|
)
|
|
from rich.table import Table
|
|
|
|
console = Console()
|
|
|
|
|
|
class RunPodVolumeUploader:
|
|
"""Uploads files to RunPod network volume via S3 API."""
|
|
|
|
def __init__(self, env_file: str = ".env.runpod"):
|
|
"""Initialize uploader with credentials from env file."""
|
|
# Load credentials
|
|
env_path = Path(__file__).parent.parent / env_file
|
|
if not env_path.exists():
|
|
console.print(f"[red]✗ Error: {env_file} not found at {env_path}[/red]")
|
|
console.print(f"[yellow]Run this script from the project root directory[/yellow]")
|
|
sys.exit(1)
|
|
|
|
load_dotenv(env_path)
|
|
|
|
# Validate required env vars
|
|
required_vars = [
|
|
"RUNPOD_S3_ACCESS_KEY",
|
|
"RUNPOD_S3_SECRET",
|
|
"RUNPOD_S3_ENDPOINT",
|
|
"RUNPOD_S3_REGION",
|
|
"RUNPOD_VOLUME_ID",
|
|
]
|
|
missing = [var for var in required_vars if not os.getenv(var)]
|
|
if missing:
|
|
console.print(f"[red]✗ Missing required env vars: {', '.join(missing)}[/red]")
|
|
sys.exit(1)
|
|
|
|
# Initialize S3 client
|
|
self.bucket = os.getenv("RUNPOD_VOLUME_ID")
|
|
self.s3_client = boto3.client(
|
|
"s3",
|
|
aws_access_key_id=os.getenv("RUNPOD_S3_ACCESS_KEY"),
|
|
aws_secret_access_key=os.getenv("RUNPOD_S3_SECRET"),
|
|
region_name=os.getenv("RUNPOD_S3_REGION"),
|
|
endpoint_url=os.getenv("RUNPOD_S3_ENDPOINT"),
|
|
config=Config(
|
|
signature_version="s3v4",
|
|
s3={"addressing_style": "path"},
|
|
retries={"max_attempts": 3, "mode": "standard"},
|
|
),
|
|
)
|
|
|
|
console.print(f"[green]✓ Initialized S3 client[/green]")
|
|
console.print(f" Endpoint: {os.getenv('RUNPOD_S3_ENDPOINT')}")
|
|
console.print(f" Volume ID: {self.bucket}")
|
|
|
|
def compute_md5(self, file_path: Path) -> str:
|
|
"""Compute MD5 checksum of local file."""
|
|
md5 = hashlib.md5()
|
|
with open(file_path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(8192), b""):
|
|
md5.update(chunk)
|
|
return md5.hexdigest()
|
|
|
|
def needs_upload(self, local_file: Path, s3_key: str) -> Tuple[bool, Optional[str]]:
|
|
"""
|
|
Check if file needs upload based on checksum.
|
|
|
|
Returns:
|
|
(needs_upload, reason)
|
|
"""
|
|
local_md5 = self.compute_md5(local_file)
|
|
|
|
try:
|
|
response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
|
|
remote_etag = response["ETag"].strip('"')
|
|
|
|
if local_md5 == remote_etag:
|
|
return False, "up-to-date (checksum match)"
|
|
else:
|
|
return True, "changed (checksum mismatch)"
|
|
except ClientError as e:
|
|
if e.response["Error"]["Code"] == "404":
|
|
return True, "new file (doesn't exist on S3)"
|
|
else:
|
|
return True, f"error checking: {e}"
|
|
|
|
def upload_file_with_progress(
|
|
self, local_file: Path, s3_key: str, progress: Progress, task: TaskID
|
|
) -> bool:
|
|
"""Upload file with progress bar."""
|
|
file_size = local_file.stat().st_size
|
|
|
|
try:
|
|
# Upload with progress callback
|
|
def callback(bytes_transferred):
|
|
progress.update(task, completed=bytes_transferred)
|
|
|
|
self.s3_client.upload_file(
|
|
str(local_file),
|
|
self.bucket,
|
|
s3_key,
|
|
Callback=callback,
|
|
ExtraArgs={"ContentType": "application/octet-stream"},
|
|
)
|
|
|
|
# Verify upload
|
|
response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
|
|
remote_size = response["ContentLength"]
|
|
|
|
if remote_size == file_size:
|
|
return True
|
|
else:
|
|
console.print(
|
|
f"[red]✗ Size mismatch: local={file_size}, remote={remote_size}[/red]"
|
|
)
|
|
self.s3_client.delete_object(Bucket=self.bucket, Key=s3_key)
|
|
return False
|
|
|
|
except ClientError as e:
|
|
console.print(f"[red]✗ Upload failed: {e}[/red]")
|
|
return False
|
|
|
|
def list_volume_contents(self, prefix: str = "") -> List[Dict]:
|
|
"""List contents of volume directory."""
|
|
try:
|
|
response = self.s3_client.list_objects_v2(Bucket=self.bucket, Prefix=prefix)
|
|
return response.get("Contents", [])
|
|
except ClientError as e:
|
|
console.print(f"[red]✗ Failed to list volume: {e}[/red]")
|
|
return []
|
|
|
|
def create_directory(self, s3_key: str, dry_run: bool = False) -> bool:
|
|
"""Create directory on volume (by uploading empty marker object)."""
|
|
if dry_run:
|
|
console.print(f"[dim]Would create directory: {s3_key}[/dim]")
|
|
return True
|
|
|
|
try:
|
|
# Check if directory already exists
|
|
objects = self.list_volume_contents(prefix=s3_key)
|
|
if objects:
|
|
console.print(f"[dim] Directory exists: {s3_key}[/dim]")
|
|
return True
|
|
|
|
# Create empty directory marker (trailing slash)
|
|
self.s3_client.put_object(Bucket=self.bucket, Key=s3_key)
|
|
console.print(f"[green]✓ Created directory: {s3_key}[/green]")
|
|
return True
|
|
except ClientError as e:
|
|
console.print(f"[red]✗ Failed to create directory {s3_key}: {e}[/red]")
|
|
return False
|
|
|
|
def upload_binaries(self, force: bool = False, dry_run: bool = False) -> bool:
|
|
"""Upload training binaries to volume."""
|
|
console.print("\n[bold cyan]Step 1: Upload Binaries[/bold cyan]")
|
|
|
|
binaries = [
|
|
"train_tft_parquet",
|
|
"train_mamba2_parquet",
|
|
"train_dqn",
|
|
"train_ppo",
|
|
]
|
|
|
|
binaries_dir = Path("target/release/examples")
|
|
if not binaries_dir.exists():
|
|
console.print(f"[red]✗ Binaries directory not found: {binaries_dir}[/red]")
|
|
console.print("[yellow]Run: cargo build --release --examples[/yellow]")
|
|
return False
|
|
|
|
# Check which binaries exist locally
|
|
existing_binaries = []
|
|
for binary in binaries:
|
|
binary_path = binaries_dir / binary
|
|
if binary_path.exists():
|
|
existing_binaries.append(binary)
|
|
else:
|
|
console.print(f"[yellow]⚠ Binary not found: {binary}[/yellow]")
|
|
|
|
if not existing_binaries:
|
|
console.print("[red]✗ No binaries found to upload[/red]")
|
|
return False
|
|
|
|
console.print(f"Found {len(existing_binaries)}/{len(binaries)} binaries")
|
|
|
|
# Upload each binary
|
|
stats = {"uploaded": 0, "skipped": 0, "failed": 0, "total_bytes": 0}
|
|
|
|
with Progress(
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
DownloadColumn(),
|
|
TransferSpeedColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as progress:
|
|
for binary in existing_binaries:
|
|
local_file = binaries_dir / binary
|
|
s3_key = f"binaries/{binary}"
|
|
file_size = local_file.stat().st_size
|
|
|
|
# Check if upload needed
|
|
if not force:
|
|
needs_upload, reason = self.needs_upload(local_file, s3_key)
|
|
if not needs_upload:
|
|
console.print(f"[green]✓ {binary}[/green]: {reason}")
|
|
stats["skipped"] += 1
|
|
continue
|
|
else:
|
|
console.print(f"[cyan]↻ {binary}[/cyan]: {reason}")
|
|
|
|
if dry_run:
|
|
console.print(f"[dim]Would upload: {binary} ({file_size:,} bytes)[/dim]")
|
|
stats["uploaded"] += 1
|
|
stats["total_bytes"] += file_size
|
|
continue
|
|
|
|
# Upload with progress bar
|
|
task = progress.add_task(f"Uploading {binary}", total=file_size)
|
|
success = self.upload_file_with_progress(local_file, s3_key, progress, task)
|
|
|
|
if success:
|
|
console.print(f"[green]✓ Uploaded: {binary} ({file_size:,} bytes)[/green]")
|
|
stats["uploaded"] += 1
|
|
stats["total_bytes"] += file_size
|
|
else:
|
|
console.print(f"[red]✗ Failed: {binary}[/red]")
|
|
stats["failed"] += 1
|
|
|
|
# Summary
|
|
console.print(f"\n[bold]Summary:[/bold]")
|
|
console.print(f" Uploaded: {stats['uploaded']}")
|
|
console.print(f" Skipped: {stats['skipped']}")
|
|
console.print(f" Failed: {stats['failed']}")
|
|
console.print(f" Total bytes: {stats['total_bytes']:,}")
|
|
|
|
return stats["failed"] == 0
|
|
|
|
def upload_test_data(self, force: bool = False, dry_run: bool = False) -> bool:
|
|
"""Upload test data files to volume."""
|
|
console.print("\n[bold cyan]Step 2: Upload Test Data[/bold cyan]")
|
|
|
|
test_data_dir = Path("test_data")
|
|
if not test_data_dir.exists():
|
|
console.print(f"[red]✗ Test data directory not found: {test_data_dir}[/red]")
|
|
return False
|
|
|
|
# Find all parquet files
|
|
parquet_files = list(test_data_dir.rglob("*.parquet"))
|
|
if not parquet_files:
|
|
console.print("[yellow]⚠ No parquet files found[/yellow]")
|
|
return True
|
|
|
|
console.print(f"Found {len(parquet_files)} parquet files")
|
|
|
|
# Upload each file
|
|
stats = {"uploaded": 0, "skipped": 0, "failed": 0, "total_bytes": 0}
|
|
|
|
with Progress(
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
DownloadColumn(),
|
|
TransferSpeedColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as progress:
|
|
for local_file in parquet_files:
|
|
# Preserve directory structure
|
|
rel_path = local_file.relative_to(test_data_dir)
|
|
s3_key = f"test_data/{rel_path}"
|
|
file_size = local_file.stat().st_size
|
|
|
|
# Check if upload needed
|
|
if not force:
|
|
needs_upload, reason = self.needs_upload(local_file, s3_key)
|
|
if not needs_upload:
|
|
console.print(f"[green]✓ {rel_path}[/green]: {reason}")
|
|
stats["skipped"] += 1
|
|
continue
|
|
else:
|
|
console.print(f"[cyan]↻ {rel_path}[/cyan]: {reason}")
|
|
|
|
if dry_run:
|
|
console.print(f"[dim]Would upload: {rel_path} ({file_size:,} bytes)[/dim]")
|
|
stats["uploaded"] += 1
|
|
stats["total_bytes"] += file_size
|
|
continue
|
|
|
|
# Upload with progress bar
|
|
task = progress.add_task(f"Uploading {rel_path}", total=file_size)
|
|
success = self.upload_file_with_progress(local_file, s3_key, progress, task)
|
|
|
|
if success:
|
|
console.print(f"[green]✓ Uploaded: {rel_path} ({file_size:,} bytes)[/green]")
|
|
stats["uploaded"] += 1
|
|
stats["total_bytes"] += file_size
|
|
else:
|
|
console.print(f"[red]✗ Failed: {rel_path}[/red]")
|
|
stats["failed"] += 1
|
|
|
|
# Summary
|
|
console.print(f"\n[bold]Summary:[/bold]")
|
|
console.print(f" Uploaded: {stats['uploaded']}")
|
|
console.print(f" Skipped: {stats['skipped']}")
|
|
console.print(f" Failed: {stats['failed']}")
|
|
console.print(f" Total bytes: {stats['total_bytes']:,}")
|
|
|
|
return stats["failed"] == 0
|
|
|
|
def create_output_directories(self, dry_run: bool = False) -> bool:
|
|
"""Create output directories on volume."""
|
|
console.print("\n[bold cyan]Step 3: Create Output Directories[/bold cyan]")
|
|
|
|
directories = [
|
|
"models/",
|
|
"models/tft/",
|
|
"models/mamba2/",
|
|
"models/dqn/",
|
|
"models/ppo/",
|
|
"logs/",
|
|
"checkpoints/",
|
|
]
|
|
|
|
# Directory creation is non-critical - models will create subdirs as needed
|
|
for directory in directories:
|
|
self.create_directory(directory, dry_run)
|
|
|
|
# Always return success - directory creation is optional
|
|
return True
|
|
|
|
def verify_upload(self) -> None:
|
|
"""Verify all uploads by listing volume contents."""
|
|
console.print("\n[bold cyan]Step 4: Verify Upload[/bold cyan]")
|
|
|
|
# List binaries
|
|
binaries = self.list_volume_contents(prefix="binaries/")
|
|
console.print(f"\n[bold]Binaries:[/bold] ({len(binaries)} files)")
|
|
for obj in binaries[:10]: # Show first 10
|
|
size_mb = obj["Size"] / (1024 * 1024)
|
|
console.print(f" {obj['Key']} ({size_mb:.1f} MB)")
|
|
|
|
# List test data
|
|
test_data = self.list_volume_contents(prefix="test_data/")
|
|
console.print(f"\n[bold]Test Data:[/bold] ({len(test_data)} files)")
|
|
for obj in test_data[:10]: # Show first 10
|
|
size_mb = obj["Size"] / (1024 * 1024)
|
|
console.print(f" {obj['Key']} ({size_mb:.1f} MB)")
|
|
|
|
# List models directory
|
|
models = self.list_volume_contents(prefix="models/")
|
|
console.print(f"\n[bold]Models:[/bold] ({len(models)} files)")
|
|
if models:
|
|
for obj in models[:5]:
|
|
size_mb = obj["Size"] / (1024 * 1024)
|
|
console.print(f" {obj['Key']} ({size_mb:.1f} MB)")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Upload binaries and test data to RunPod network volume via S3 API",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Upload binaries only
|
|
./scripts/upload_to_runpod_volume.py --binaries
|
|
|
|
# Upload test data only
|
|
./scripts/upload_to_runpod_volume.py --test-data
|
|
|
|
# Upload everything
|
|
./scripts/upload_to_runpod_volume.py --all
|
|
|
|
# Force re-upload (ignore checksums)
|
|
./scripts/upload_to_runpod_volume.py --all --force
|
|
|
|
# Dry run (show what would be uploaded)
|
|
./scripts/upload_to_runpod_volume.py --all --dry-run
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--binaries", action="store_true", help="Upload training binaries"
|
|
)
|
|
parser.add_argument(
|
|
"--test-data", action="store_true", help="Upload test data files"
|
|
)
|
|
parser.add_argument(
|
|
"--all", action="store_true", help="Upload binaries and test data"
|
|
)
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Force re-upload (ignore checksums)",
|
|
)
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Show what would be uploaded without uploading",
|
|
)
|
|
parser.add_argument(
|
|
"--env-file",
|
|
default=".env.runpod",
|
|
help="Path to env file (default: .env.runpod)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Validate arguments
|
|
if not (args.binaries or args.test_data or args.all):
|
|
parser.print_help()
|
|
console.print("\n[red]Error: Must specify --binaries, --test-data, or --all[/red]")
|
|
sys.exit(1)
|
|
|
|
# Initialize uploader
|
|
uploader = RunPodVolumeUploader(env_file=args.env_file)
|
|
|
|
# Execute uploads
|
|
success = True
|
|
|
|
if args.all or args.binaries:
|
|
if not uploader.upload_binaries(force=args.force, dry_run=args.dry_run):
|
|
success = False
|
|
|
|
if args.all or args.test_data:
|
|
if not uploader.upload_test_data(force=args.force, dry_run=args.dry_run):
|
|
success = False
|
|
|
|
if args.all:
|
|
if not uploader.create_output_directories(dry_run=args.dry_run):
|
|
success = False
|
|
|
|
# Verify upload
|
|
if not args.dry_run:
|
|
uploader.verify_upload()
|
|
|
|
# Exit with appropriate status
|
|
if success:
|
|
console.print("\n[bold green]✓ Upload complete![/bold green]")
|
|
sys.exit(0)
|
|
else:
|
|
console.print("\n[bold red]✗ Upload failed![/bold red]")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|