Files
dataset-uploader/scripts/download_with_resume.py

130 lines
4.8 KiB
Python

#!/usr/bin/env python3
"""Download large files with automatic resume capability."""
import sys
import time
from pathlib import Path
import httpx
def download_with_resume(url: str, output_path: Path, max_retries: int = 5):
"""Download a file with automatic resume on failure."""
output_path.parent.mkdir(parents=True, exist_ok=True)
for attempt in range(max_retries):
try:
# Check if partial download exists
start_byte = 0
if output_path.exists():
start_byte = output_path.stat().st_size
print(
f"Resuming from byte {start_byte} ({start_byte / (1024**2):.2f} MB)"
)
mode = "ab"
else:
mode = "wb"
# Set up headers for range request
headers = {}
if start_byte > 0:
headers["Range"] = f"bytes={start_byte}-"
# Download with timeout and retries
with (
httpx.Client(timeout=httpx.Timeout(60.0, read=30.0)) as client,
client.stream("GET", url, headers=headers) as response,
):
# Check if server supports resume
if start_byte > 0 and response.status_code != 206:
print("Server doesn't support resume, restarting download")
output_path.unlink()
start_byte = 0
response = client.stream("GET", url).__enter__()
response.raise_for_status()
# Get total size
if response.status_code == 206: # Partial content
content_range = response.headers.get("content-range", "")
if content_range:
total_size = int(content_range.split("/")[-1])
else:
total_size = (
int(response.headers.get("content-length", 0))
+ start_byte
)
else:
total_size = int(response.headers.get("content-length", 0))
print(f"Total size: {total_size / (1024**2):.2f} MB")
# Download with progress
with output_path.open(mode) as f:
downloaded = start_byte
last_update = time.time()
for chunk in response.iter_bytes(chunk_size=1024 * 1024):
f.write(chunk)
downloaded += len(chunk)
# Update progress every second
if time.time() - last_update > 1:
progress = (
(downloaded / total_size * 100) if total_size else 0
)
sys.stdout.write(
f"\rProgress: "
f"{downloaded / (1024**2):.2f}/"
f"{total_size / (1024**2):.2f} MB "
f"({progress:.1f}%)"
)
sys.stdout.flush()
last_update = time.time()
# Final update
progress = (downloaded / total_size * 100) if total_size else 0
print(
f"\rProgress: "
f"{downloaded / (1024**2):.2f}/"
f"{total_size / (1024**2):.2f} MB "
f"({progress:.1f}%)"
)
print("Download complete!")
return True
except (
httpx.TimeoutException,
httpx.NetworkError,
httpx.RemoteProtocolError,
) as e:
print(f"\nDownload interrupted: {e}")
if attempt < max_retries - 1:
print(f"Retrying in 5 seconds... (attempt {attempt + 2}/{max_retries})")
time.sleep(5)
else:
print("Max retries reached. Download failed.")
return False
except KeyboardInterrupt:
print("\nDownload cancelled by user")
return False
except Exception as e:
print(f"\nUnexpected error: {e}")
return False
return False
if __name__ == "__main__":
if len(sys.argv) != 3:
print("Usage: python download_with_resume.py <url> <output_path>")
sys.exit(1)
url = sys.argv[1]
output_path = Path(sys.argv[2])
success = download_with_resume(url, output_path)
sys.exit(0 if success else 1)