Files
dataset-uploader/scripts/download_with_progress.py

119 lines
3.9 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Wrapper for rdf_dataset_downloader.py that adds progress output for parallel processing.
This wrapper monitors the actual download progress and outputs PROGRESS markers.
"""
import re
import subprocess
import sys
import threading
import time
from pathlib import Path
def monitor_download_output(proc):
"""Monitor download output and emit progress markers."""
# Patterns to detect download progress
progress_patterns = [
# httpx/requests download progress patterns
(r"(\d+)%\|", lambda m: float(m.group(1))), # 45%|████
(
r"(\d+\.?\d*)\s*/\s*(\d+\.?\d*)\s*MB",
lambda m: (float(m.group(1)) / float(m.group(2))) * 100,
), # 45.5/100 MB
(
r"(\d+\.?\d*)\s*MB\s*/\s*(\d+\.?\d*)\s*MB",
lambda m: (float(m.group(1)) / float(m.group(2))) * 100,
), # 45.5 MB / 100 MB
(
r"Downloaded\s+(\d+)",
lambda m: min(10, float(m.group(1)) / 1000),
), # Downloaded bytes
(r"(\d+\.?\d*)%", lambda m: float(m.group(1))), # Simple percentage
# Stage detection
(r"Downloading.*from", lambda m: 5), # Starting download
(r"Connecting to", lambda m: 2), # Connecting
(r"Extracting", lambda m: 85), # Extracting archive
(r"Extracted|Complete|Done|Saved", lambda m: 100), # Complete
]
last_progress = 0
last_emit_time = 0
# Since downloads often don't show progress, we'll simulate it based on time
download_start = time.time()
estimated_duration = 30 # Estimate 30 seconds for download
def emit_time_based_progress():
"""Emit progress based on elapsed time if no real progress detected."""
nonlocal last_progress, last_emit_time
while proc.poll() is None:
time.sleep(1)
elapsed = time.time() - download_start
# Estimate progress based on time (max 80% to leave room for extraction)
time_progress = min(80, (elapsed / estimated_duration) * 80)
if time_progress > last_progress + 5: # Only update if significant change
last_progress = time_progress
print(f"\nPROGRESS: {time_progress:.0f}%", flush=True)
last_emit_time = time.time()
# Start time-based progress thread
progress_thread = threading.Thread(target=emit_time_based_progress, daemon=True)
progress_thread.start()
# Monitor actual output
for line in iter(proc.stdout.readline, ""):
if not line:
break
# Pass through original output
print(line, end="", flush=True)
# Try to extract real progress
for pattern, extractor in progress_patterns:
match = re.search(pattern, line, re.IGNORECASE)
if match:
try:
progress = extractor(match)
current_time = time.time()
if (
progress > last_progress
and (current_time - last_emit_time) > 0.5
):
last_progress = progress
last_emit_time = current_time
print(f"\nPROGRESS: {progress:.0f}%", flush=True)
break
except Exception:
pass
# Ensure we reach 100% when done
if proc.wait() == 0 and last_progress < 100:
print("\nPROGRESS: 100%", flush=True)
def main():
"""Run the downloader with progress monitoring."""
# Run the actual downloader
proc = subprocess.Popen(
[
sys.executable,
"-u",
str(Path(__file__).parent / "rdf_dataset_downloader.py"),
*sys.argv[1:],
],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
monitor_download_output(proc)
return proc.returncode
if __name__ == "__main__":
sys.exit(main())