lars/aiRtrafficNN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use async for data downloading

Browse Source
This commit is contained in:
Lars Bogner
2025-11-12 14:23:31 +01:00
parent bc6d436b85
commit 1212c2deac
7 changed files with 735 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Code/python/src/adsbpy/fetch.py
+276 -63
View File
@@ -1,96 +1,309 @@
#!/usr/bin/env python3
import os
import subprocess
import urllib.request
from pathlib import Path
import asyncio
import aiohttp
import aiofiles
import tarfile
import shutil
from pathlib import Path
from typing import List, Set
import logging
from concurrent.futures import ThreadPoolExecutor
import os
# --- configuration ---
BASE_DIR = Path(__file__).parent.resolve() / Path("../../../..") # root of project
DATA_DIR = BASE_DIR / Path("data") # where extracted tars go
CSV_DIR = BASE_DIR / Path("data/csv") # where final csvs go
EXTRACT_ADSB = BASE_DIR / Path("Code/cpp/extract-adsb") # external binary
BASE_DIR = Path(__file__).parent.resolve() / Path("../../../..")
DATA_DIR = BASE_DIR / "data"
CSV_DIR = BASE_DIR / "data/csv"
EXTRACT_ADSB = BASE_DIR / "Code/cpp/extract-adsb"
# Performance tuning
MAX_CONCURRENT_DOWNLOADS = 2 # Limit concurrent HTTP requests
MAX_CONCURRENT_EXTRACTIONS = 4 # Limit concurrent tar extractions (CPU intensive)
MAX_CONCURRENT_DAYS = 3 # Limit concurrent day processing
# -----------------------
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def download(url: str, dest: Path):
dest.parent.mkdir(parents=True, exist_ok=True)
print(f"Downloading {url} → {dest}")
with urllib.request.urlopen(url) as r, open(dest, "wb") as f:
shutil.copyfileobj(r, f)
def pretty_path(path: os.PathLike | str) -> Path:
"""Return a path relative to the current working directory."""
path = Path(path).resolve()
return Path(os.path.relpath(path, Path.cwd()))
def pretty_url(url: str) -> str:
"""Return a shortened version of the URL for logging."""
if len(url) <= 60:
return url
return f"{url[:30]}...{url[-30:]}"
async def download(session: aiohttp.ClientSession, url: str, dest: Path, semaphore: asyncio.Semaphore) -> None:
"""Download a file asynchronously with semaphore limiting and retry logic."""
async with semaphore:
dest.parent.mkdir(parents=True, exist_ok=True)
logger.info(f"Downloading {pretty_url(url)} → {pretty_path(dest)}")
max_retries = 3
retry_delay = 10 # seconds
timeout = 600 # seconds
for attempt in range(max_retries + 1):
try:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status == 429: # Too Many Requests
wait_time = int(response.headers.get('Retry-After', retry_delay * (attempt + 1)))
logger.warning(f"Rate limited for {pretty_url(url)}. Waiting {wait_time}s before retry {attempt + 1}/{max_retries}")
await asyncio.sleep(wait_time)
continue
response.raise_for_status()
async with aiofiles.open(dest, 'wb') as f:
async for chunk in response.content.iter_chunked(8192 * 4):
await f.write(chunk)
logger.info(f"Download completed: {pretty_path(dest)}")
return # Success, exit function
except aiohttp.ClientResponseError as e:
if e.status >= 500 and attempt < max_retries: # Server errors are retryable
logger.warning(f"Server error {e.status} for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
await asyncio.sleep(retry_delay * (2 ** attempt)) # Exponential backoff
continue
else:
logger.error(f"HTTP error {e.status} for {pretty_url(url)}: {e}")
raise
except aiohttp.ClientConnectorError as e:
if attempt < max_retries:
logger.warning(f"Connection error for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
await asyncio.sleep(retry_delay * (2 ** attempt))
continue
else:
logger.error(f"Connection failed for {pretty_url(url)}: {e}")
raise
except asyncio.TimeoutError as e:
if attempt < max_retries:
logger.warning(f"Timeout for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
await asyncio.sleep(retry_delay * (2 ** attempt))
continue
else:
logger.error(f"Download timeout for {pretty_url(url)}: {e}")
raise
except Exception as e:
if attempt < max_retries:
logger.warning(f"Unexpected error for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
await asyncio.sleep(retry_delay * (2 ** attempt))
continue
else:
logger.error(f"Download failed for {pretty_url(url)} after {max_retries} attempts: {e}")
raise
# This should never be reached, but just in case
raise Exception(f"Download failed for {pretty_url(url)} after {max_retries} attempts")
def get_name_from_url(url: str) -> str:
"""Extract filename from URL."""
return url.split("/")[-1]
def extract_adsb_data(cmd):
print(f"Running: {' '.join(map(str, cmd))}")
subprocess.run(cmd, check=True)
def extract_tar(tar_paths: list[Path], extract_dir: Path):
async def extract_adsb_data(cmd: List[str]) -> None:
"""Run extract-adsb binary asynchronously."""
logger.info(f"Running: {' '.join(cmd)}")
process = await asyncio.create_subprocess_exec(*cmd)
await process.wait()
if process.returncode != 0:
raise Exception(f"Command failed with return code {process.returncode}: {' '.join(cmd)}")
def extract_tar_sync(tar_paths: List[Path], extract_dir: Path) -> Path:
"""Synchronous tar extraction (CPU-bound, run in thread pool)."""
extract_dir.mkdir(parents=True, exist_ok=True)
if len(tar_paths) == 1:
tar_path = tar_paths[0]
else:
# concatenate parts
# Remove all files without .tar.aa, .tar.ab, ... suffixes
# Concatenate multi-part tar files
tar_paths = [p for p in tar_paths if len(p.suffix) == 3]
tar_path = tar_paths[0].stem # files are named foo.tar.aa, foo.tar.ab, ...
tar_path = extract_dir / tar_path
print(f"Concatenating {len(tar_paths)} parts → {tar_path}")
with open(tar_path, "wb") as out:
tar_path = extract_dir / tar_paths[0].stem
logger.info(f"Concatenating {len(tar_paths)} parts → {pretty_path(tar_path)}")
with open(tar_path, 'wb') as outfile:
for part in tar_paths:
with open(part, "rb") as pf:
shutil.copyfileobj(pf, out)
print(f"Extracting {tar_path} → {extract_dir}")
with open(part, 'rb') as infile:
shutil.copyfileobj(infile, outfile)
logger.info(f"Extracting {pretty_path(tar_path)} → {pretty_path(extract_dir)}")
with tarfile.open(tar_path) as tf:
tf.extractall(path=extract_dir)
return extract_dir
def move_csv_files(src_dir: Path, dest_dir: Path):
async def extract_tar(tar_paths: List[Path], extract_dir: Path, executor: ThreadPoolExecutor) -> Path:
"""Run tar extraction in thread pool to avoid blocking event loop."""
loop = asyncio.get_event_loop()
return await loop.run_in_executor(executor, extract_tar_sync, tar_paths, extract_dir)
async def move_csv_files(src_dir: Path, dest_dir: Path) -> None:
"""Move all CSV files from source to destination directory."""
dest_dir.mkdir(parents=True, exist_ok=True)
for p in src_dir.rglob("*.csv"):
dest = dest_dir / p.name
shutil.move(str(p), dest)
shutil.move(str(p), str(dest))
def remove_dir(dir_path: Path):
for p in dir_path.rglob("*"):
if p.is_file():
p.unlink()
for p in dir_path.rglob("*"):
if p.is_dir():
if not any(p.iterdir()):
p.rmdir()
else:
remove_dir(p)
dir_path.rmdir()
def fetch_lines(lines: list[str]):
for line in lines:
parts = line.strip().split(",")
if not parts:
continue
tar_files = []
for url in parts:
filename = DATA_DIR / get_name_from_url(url)
tar_files.append(filename)
download(url, filename)
async def remove_dir(dir_path: Path) -> None:
"""Remove directory and all its contents."""
if dir_path.exists():
shutil.rmtree(dir_path)
async def remove_file(file_path: Path) -> None:
"""Remove a single file."""
if file_path.exists():
file_path.unlink()
async def download_day_parts(session: aiohttp.ClientSession, urls: List[str],
download_semaphore: asyncio.Semaphore) -> List[Path]:
"""Download all parts for a single day concurrently."""
tar_files = []
download_tasks = []
for url in urls:
filename = DATA_DIR / get_name_from_url(url)
tar_files.append(filename)
download_tasks.append(
download(session, url, filename, download_semaphore)
)
await asyncio.gather(*download_tasks)
return tar_files
async def process_single_day(session: aiohttp.ClientSession, line: str,
download_semaphore: asyncio.Semaphore,
extraction_executor: ThreadPoolExecutor) -> None:
"""Process a single day's data with concurrent downloads and parallel extraction."""
parts = line.strip().split(",")
if not parts or not parts[0]:
return
if len(parts) > 1:
# Remove any part that is not of ending .tar.aa, .tar.ab, etc.
parts = [p for p in parts if len(p.split(".")[-1]) == 2]
logger.info(f"Processing day with {len(parts)} parts")
try:
# Download all parts for this day concurrently
tar_files = await download_day_parts(session, parts, download_semaphore)
# Extract the tar files
extract_dir = tar_files[0].parent / tar_files[0].stem
extract_tar(tar_files, extract_dir)
await extract_tar(tar_files, extract_dir, extraction_executor)
# Process ADSB data
traces_dir = extract_dir / "traces"
if traces_dir.exists():
extract_adsb_data([EXTRACT_ADSB, str(traces_dir)])
await extract_adsb_data([str(EXTRACT_ADSB), str(traces_dir)])
# Move CSV files
csv_files_dir = CSV_DIR / extract_dir.name
move_csv_files(extract_dir, csv_files_dir)
# clean up
for tar_file in tar_files:
if tar_file.exists():
tar_file.unlink()
remove_dir(extract_dir)
await move_csv_files(extract_dir, csv_files_dir)
# Cleanup
cleanup_tasks = [
asyncio.create_task(remove_file(tar_file))
for tar_file in tar_files
if tar_file.exists()
]
cleanup_tasks.append(asyncio.create_task(remove_dir(extract_dir)))
await asyncio.gather(*cleanup_tasks)
logger.info(f"Completed processing day: {extract_dir.name}")
except Exception as e:
logger.error(f"Failed to process day: {e}")
raise
else:
# If successful, write to ignore file
await write_line_to_ignore(BASE_DIR / "Code/python/adsbpy/.ignore", line.strip())
def fetch_file(file: Path):
with open(file) as f:
lines = f.readlines()
fetch_lines(lines)
async def write_line_to_ignore(file: Path, line: str) -> None:
"""Append a line to the ignore file asynchronously."""
async with aiofiles.open(file, 'a') as f:
await f.write(line + '\n')
async def process_days_concurrently(lines: List[str]) -> None:
"""Process multiple days concurrently with controlled parallelism."""
# Create semaphores for controlling concurrency
download_semaphore = asyncio.Semaphore(MAX_CONCURRENT_DOWNLOADS)
day_semaphore = asyncio.Semaphore(MAX_CONCURRENT_DAYS)
# Thread pool for CPU-bound tar extraction
with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_EXTRACTIONS) as extraction_executor:
connector = aiohttp.TCPConnector(limit=MAX_CONCURRENT_DOWNLOADS)
async with aiohttp.ClientSession(connector=connector) as session:
async def process_day_with_semaphore(line: str) -> None:
"""Wrapper to limit concurrent day processing."""
async with day_semaphore:
await process_single_day(session, line, download_semaphore, extraction_executor)
# Process all days concurrently with limits
tasks = [
process_day_with_semaphore(line)
for line in lines
if line.strip()
]
# Process in batches to avoid overwhelming the system
batch_size = MAX_CONCURRENT_DAYS * 5
for i in range(0, len(tasks), batch_size):
batch = tasks[i:i + batch_size]
logger.info(f"Processing batch {i//batch_size + 1}/{(len(tasks)-1)//batch_size + 1}")
await asyncio.gather(*batch, return_exceptions=True)
async def fetch_lines(lines: List[str]) -> None:
"""Process multiple lines (days) concurrently."""
await process_days_concurrently(lines)
async def fetch_file(file: Path) -> None:
"""Fetch and process URLs from a file."""
async with aiofiles.open(file, 'r') as f:
lines = await f.readlines()
logger.info(f"Processing {len(lines)} days")
await fetch_lines(lines)
async def main() -> None:
"""Main entry point if running as a script."""
import sys
if len(sys.argv) > 1:
input_file = Path(sys.argv[1])
else:
input_file = Path("urls.txt")
if input_file.exists():
logger.info(f"Starting processing of {input_file}")
await fetch_file(input_file)
logger.info("Processing completed successfully")
else:
logger.error(f"Input file {input_file} not found")
sys.exit(1)
if __name__ == "__main__":
asyncio.run(main())