Files
dataset-uploader/scripts/convert_rdf_streaming_progress.py

106 lines
4.0 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Wrapper for convert_rdf_to_hf_dataset_streaming_parallel.py that adds progress output
for subprocess monitoring in parallel processing.
"""
import subprocess
import sys
import re
from pathlib import Path
import threading
import time
def monitor_and_output_progress(proc):
"""Monitor subprocess output and emit simple progress markers."""
progress_patterns = [
# Streaming converter specific patterns
(
r"Chunk (\d+)/(\d+)",
lambda m: (float(m.group(1)) / float(m.group(2))) * 70 if m.group(2) != "0" else 0,
), # Chunk progress
(r"Processing chunk (\d+)", lambda m: min(5 + int(m.group(1)) * 5, 70)), # Processing chunks
(r"Worker \d+: bytes", lambda m: 10), # Worker assignment
(r"✓ File divided", lambda m: 15), # File divided
(r"Parsing chunks", lambda m: 20), # Starting parsing
(r"Processed (\d+) triples", lambda m: min(70, 20 + int(m.group(1)) / 1000)), # Triple count
# Phase detection with incremental progress
(r"Phase 1:", lambda m: 5), # Starting phase 1
(r"Dividing file", lambda m: 10), # Dividing file
(r"Phase 2:", lambda m: 20), # Starting phase 2
(r"Phase 3:", lambda m: 75), # Starting phase 3
(r"Creating dataset", lambda m: 80), # Creating dataset
(r"Writing to Parquet", lambda m: 85), # Writing parquet
(r"Saving dataset", lambda m: 90), # Saving dataset
(r"Building dataset", lambda m: 92), # Building dataset
(r"Dataset saved", lambda m: 95), # Dataset saved
(r"✓.*Successfully|Complete|Done", lambda m: 100), # Complete markers
# Memory usage patterns (informational)
(r"Memory usage:", lambda m: None), # Just pass through
# Generic percentage patterns
(r"(\d+)%", lambda m: float(m.group(1))), # Direct percentage
(r"\[(\d+)%\]", lambda m: float(m.group(1))), # Bracketed percentage
(r"Progress: (\d+\.?\d*)%", lambda m: float(m.group(1))), # Progress: X%
]
last_progress = 0
last_emit_time = 0
MIN_EMIT_INTERVAL = 0.5 # Minimum time between progress emissions
# Track if we're in streaming mode
streaming_mode = False
for line in iter(proc.stdout.readline, ""):
if not line:
break
# Pass through the original output
print(line, end="", flush=True)
# Detect streaming mode
if "streaming" in line.lower() or "chunk" in line.lower():
streaming_mode = True
# Try to extract and emit progress
for pattern, extractor in progress_patterns:
match = re.search(pattern, line, re.IGNORECASE)
if match:
try:
progress = extractor(match)
if progress is not None: # Skip None values (informational patterns)
current_time = time.time()
# Only emit if progress increased and enough time has passed
if progress > last_progress and (current_time - last_emit_time) > MIN_EMIT_INTERVAL:
last_progress = progress
last_emit_time = current_time
# Emit simple progress marker
print(f"\nPROGRESS: {progress:.0f}", flush=True)
break
except Exception:
pass
# Ensure we hit 100% on successful completion
if proc.wait() == 0 and last_progress < 100:
print("\nPROGRESS: 100", flush=True)
def main():
"""Run the streaming converter with progress monitoring."""
# Pass all arguments to the actual converter
proc = subprocess.Popen(
[sys.executable, "-u", str(Path(__file__).parent / "convert_rdf_to_hf_dataset_streaming_parallel.py")]
+ sys.argv[1:],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
monitor_and_output_progress(proc)
return proc.returncode
if __name__ == "__main__":
sys.exit(main())