Coverage for lunar_toolkit/downloader.py: 93%
248 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 02:58 +0000
1"""PdsDownloader — parallel file retrieval from the PDS archive."""
3import contextlib
4import contextvars
5import email.utils
6import hashlib
7import json
8import logging
9import os
10import random
11import threading
12import time
13import urllib.error
14import urllib.parse
15import urllib.request
16from collections.abc import Callable, Iterator
17from concurrent.futures import ThreadPoolExecutor, as_completed
18from pathlib import Path
20import pandas as pd
22from .constants import APP_NAME, APP_VERSION, DEFAULT_CACHE_DIR, modernize_pds_url
23from .tools import reuse_existing
25log = logging.getLogger(__name__)
26_USER_AGENT = f"{APP_NAME}/{APP_VERSION} (python-urllib)"
28Progress = Callable[[str, int, int | None], None]
29"""``progress(file name, bytes done, bytes total or None)``, called per chunk; raising aborts.
31While the host is paused (throttling cooldown) it is called about once a second with
32``done = -1``.
33"""
35_PROGRESS: contextvars.ContextVar[Progress | None] = contextvars.ContextVar("download_progress", default=None)
38@contextlib.contextmanager
39def report_progress(callback: Progress) -> Iterator[None]:
40 """Send the progress of every [`download_file`][lunar_toolkit.downloader.download_file] in this thread /
41 context to *callback*."""
42 token = _PROGRESS.set(callback)
43 try:
44 yield
45 finally:
46 _PROGRESS.reset(token)
49class PdsDownloader:
50 """Download and cache files from the PDS archive.
52 Parameters
53 ----------
54 cache_dir : Path
55 Base directory for the local cache (default ~/.cache/lunar_toolkit).
56 max_workers : int
57 Thread-pool size for parallel downloads (default 8).
58 """
60 def __init__(
61 self,
62 cache_dir: Path = DEFAULT_CACHE_DIR,
63 max_workers: int = 4,
64 ) -> None:
65 self.cache_dir = Path(cache_dir)
66 self.max_workers = max_workers
68 def cache_path(self, url: str) -> Path:
69 """Return the local path that would be used to cache *url*."""
70 url_hash = hashlib.md5(url.encode()).hexdigest()[:8]
71 return self.cache_dir / f"{url_hash}_{url_filename(url)}"
73 def ensure_cached(
74 self,
75 url: str,
76 force: bool = False,
77 show_progress: bool = True,
78 expected_md5: str | None = None,
79 ) -> Path:
80 """Return local path for *url*, downloading if not already cached."""
81 dest = self.cache_path(url)
82 if dest.exists() and not force: 82 ↛ 84line 82 didn't jump to line 84 because the condition on line 82 was always true
83 return dest
84 return download_file(url, dest, show_progress=show_progress, expected_md5=expected_md5)
86 def download(
87 self,
88 product_rows: pd.DataFrame,
89 output_dir: Path,
90 url_col: str = "PDS_URL",
91 show_progress: bool = True,
92 force: bool = False,
93 ) -> pd.DataFrame:
94 """Download PDS image files in parallel.
96 Parameters
97 ----------
98 product_rows : pd.DataFrame
99 Must contain a *url_col* column (e.g. from CumIndex.lookup()).
100 output_dir : Path
101 Destination directory for .IMG files.
102 url_col : str
103 Column holding the download URL.
104 show_progress : bool
105 Show per-file progress bars.
106 force : bool
107 Download again even when ``<output_dir>/<file name>`` exists
108 (by default it is reused; downloads land via ``.part`` files,
109 so an existing file is complete).
111 Returns
112 -------
113 Copy of product_rows with LOCAL_PATH and DOWNLOAD_ERROR columns added.
114 """
115 results = self._download_all(list(product_rows[url_col]), output_dir, show_progress, force)
116 paths = {u: str(path) if path else None for u, (path, _) in results.items()}
117 df = product_rows.copy()
118 df["LOCAL_PATH"] = df[url_col].map(paths.get)
119 df["DOWNLOAD_ERROR"] = df[url_col].map(lambda u: results.get(u, (None, None))[1])
120 return df
122 def download_urls(
123 self,
124 urls: list[str],
125 output_dir: Path,
126 show_progress: bool = True,
127 force: bool = False,
128 ) -> dict[str, Path | None]:
129 """Download a plain list of URLs in parallel.
131 Existing files are reused unless *force* (see
132 [`download`][lunar_toolkit.downloader.PdsDownloader.download]). Returns a dict mapping each URL to its
133 local path (or None on error).
134 """
135 path_map: dict[str, Path | None] = {}
136 for url, (path, err) in self._download_all(urls, output_dir, show_progress, force).items():
137 if err: 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
138 log.error("%s: %s", url, err)
139 path_map[url] = path
140 return path_map
142 def _download_all(
143 self, urls: list[str], output_dir: Path, show_progress: bool, force: bool
144 ) -> dict[str, tuple[Path | None, str | None]]:
145 """``url → (local path, error message)`` for parallel downloads into *output_dir*."""
146 output_dir = Path(output_dir)
147 output_dir.mkdir(parents=True, exist_ok=True)
148 results: dict[str, tuple[Path | None, str | None]] = {}
149 with ThreadPoolExecutor(max_workers=self.max_workers) as pool:
150 futures = [
151 pool.submit(_download_one, url, output_dir / url_filename(url), show_progress, force)
152 for url in urls
153 ]
154 for future in as_completed(futures):
155 url, path, err = future.result()
156 results[url] = (path, err)
157 return results
160class ChecksumError(IOError):
161 """Raised when a downloaded file does not match its expected MD5."""
164def file_md5(path: Path, chunk: int = 8 << 20) -> str:
165 """Return the hex MD5 digest of *path* (streamed, constant memory)."""
166 h = hashlib.md5()
167 with open(path, "rb") as f:
168 while True:
169 buf = f.read(chunk)
170 if not buf:
171 break
172 h.update(buf)
173 return h.hexdigest()
176def url_filename(url: str) -> str:
177 """File name at the end of *url*'s path (query string ignored)."""
178 return Path(urllib.parse.urlparse(url).path).name
181RETRY_CODES = (429, 500, 502, 503, 504)
182THROTTLE_CODES = (429, 503)
183MAX_PER_HOST_ENV = "LUNAR_MAX_PER_HOST"
184MIN_INTERVAL_ENV = "LUNAR_MIN_INTERVAL"
185MAX_WAIT_ENV = "LUNAR_MAX_WAIT"
186MAX_RETRY_AFTER = 7200.0
187COOLDOWN_FILE = DEFAULT_CACHE_DIR / "host_cooldowns.json"
190class _Host:
191 """One host's download slots, request pacing and throttling cooldown.
193 The cooldown is wall-clock time shared by every worker in this process and saved
194 to `COOLDOWN_FILE`, so a restarted run honours a server's earlier
195 ``Retry-After`` instead of tripping its rate limit again.
196 """
198 def __init__(self, name: str, limit: int, min_interval: float) -> None:
199 self.name = name
200 self.slots = threading.BoundedSemaphore(limit)
201 self.lock = threading.Lock()
202 self.min_interval = min_interval
203 self.last_start = 0.0
204 self.cooldown_until = _saved_cooldowns().get(name, 0.0)
206 def cool(self, seconds: float) -> None:
207 with self.lock:
208 until = time.time() + seconds
209 if until > self.cooldown_until: 209 ↛ exitline 209 didn't jump to the function exit
210 self.cooldown_until = until
211 _save_cooldown(self.name, until)
213 def wait(self, poll: Callable[[], None] | None = None) -> None:
214 """Sleep out the cooldown; with *poll*, in 1 s steps calling it (it may raise to abort)."""
215 remaining = self.cooldown_until - time.time()
216 if remaining > 0:
217 resume = time.strftime("%H:%M:%S", time.localtime(self.cooldown_until))
218 log.info("%s asked us to slow down; waiting %.0f s (until %s)", self.name, remaining, resume)
219 if poll is None:
220 time.sleep(remaining)
221 return
222 while remaining > 0: 222 ↛ exitline 222 didn't return from function 'wait' because the condition on line 222 was always true
223 poll()
224 time.sleep(min(1.0, remaining))
225 remaining = self.cooldown_until - time.time()
227 def pace(self) -> None:
228 with self.lock:
229 now = time.monotonic()
230 start = max(now, self.last_start + self.min_interval)
231 self.last_start = start
232 if start > now: 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 time.sleep(start - now)
236_HOSTS: dict[str, _Host] = {}
237_HOSTS_LOCK = threading.Lock()
240def _host(url: str) -> _Host:
241 name = urllib.parse.urlparse(url).netloc.lower()
242 with _HOSTS_LOCK:
243 if name not in _HOSTS:
244 limit = max(1, int(os.environ.get(MAX_PER_HOST_ENV, "3")))
245 _HOSTS[name] = _Host(name, limit, float(os.environ.get(MIN_INTERVAL_ENV, "1.0")))
246 return _HOSTS[name]
249def _saved_cooldowns() -> dict[str, float]:
250 try:
251 data = json.loads(COOLDOWN_FILE.read_text())
252 except (OSError, ValueError):
253 return {}
254 now = time.time()
255 return {str(k): float(v) for k, v in data.items() if isinstance(v, (int, float)) and v > now}
258def _save_cooldown(name: str, until: float) -> None:
259 data = _saved_cooldowns()
260 data[name] = until
261 try:
262 COOLDOWN_FILE.parent.mkdir(parents=True, exist_ok=True)
263 tmp = COOLDOWN_FILE.with_suffix(f".{os.getpid()}.tmp")
264 tmp.write_text(json.dumps(data))
265 tmp.replace(COOLDOWN_FILE)
266 except OSError as exc:
267 log.debug("Could not save host cooldowns: %s", exc)
270def download_file(
271 url: str,
272 dest: Path,
273 show_progress: bool = True,
274 expected_md5: str | None = None,
275 resume: bool = True,
276 timeout: float = 60.0,
277 retries: int = 8,
278 retry_codes: tuple[int, ...] = RETRY_CODES,
279 max_wait: float | None = None,
280 progress: Progress | None = None,
281) -> Path:
282 """Download *url* to *dest*, throttled per host and retried on throttling / errors.
284 Politeness: at most ``$LUNAR_MAX_PER_HOST`` (default 3) downloads run against one
285 host at a time, and their requests start at least ``$LUNAR_MIN_INTERVAL`` seconds
286 apart (default 1). Throttling replies (429, 503) pause every download to that host
287 for the server's ``Retry-After`` (up to 2 h; else exponential back-off) and do not
288 count as failures; they give up only after *max_wait* seconds in total (default
289 ``$LUNAR_MAX_WAIT``, else 6 h). Other *retry_codes* (5xx), connection errors and
290 interrupted transfers are retried *retries* times with back-off. Each retry resumes
291 the ``.part`` file; checksum mismatches are not retried. See `_download_once`.
293 *progress* (default: the [`report_progress`][lunar_toolkit.downloader.report_progress] callback, if any)
294 gets the bytes done after each chunk; an exception it raises stops the download and keeps the
295 ``.part`` file for a later resume.
296 """
297 progress = progress or _PROGRESS.get()
298 url = modernize_pds_url(url)
299 host = _host(url)
300 budget = max_wait if max_wait is not None else float(os.environ.get(MAX_WAIT_ENV, 6 * 3600))
301 deadline = time.time() + budget
302 failures = throttles = 0
303 poll = (lambda: progress(Path(dest).name, -1, None)) if progress else None
304 while True:
305 host.wait(poll)
306 try:
307 with host.slots:
308 host.pace()
309 return _download_once(url, dest, show_progress, expected_md5, resume, timeout, progress)
310 except urllib.error.HTTPError as exc:
311 if exc.code not in retry_codes:
312 raise
313 if exc.code in THROTTLE_CODES:
314 wait = _retry_after(exc) or _backoff(throttles)
315 throttles += 1
316 else:
317 failures += 1
318 if failures > retries:
319 raise
320 wait = _backoff(failures - 1)
321 if time.time() + wait > deadline:
322 raise
323 log.warning("HTTP %d from %s; pausing that host %.0f s", exc.code, host.name, wait)
324 except ChecksumError:
325 raise
326 except OSError as exc: # URLError, timeouts, resets, incomplete transfers
327 failures += 1
328 if failures > retries:
329 raise
330 wait = _backoff(failures - 1)
331 log.warning("%s for %s; retrying in %.0f s", exc, url, wait)
332 host.cool(wait)
335def _backoff(attempt: int) -> float:
336 return min(120.0, 2.0 * 2 ** min(attempt, 10)) * random.uniform(0.75, 1.25)
339def _retry_after(exc: urllib.error.HTTPError) -> float | None:
340 value = exc.headers.get("Retry-After") if exc.headers else None
341 if value is None:
342 return None
343 try:
344 seconds = float(value)
345 except ValueError:
346 try:
347 seconds = email.utils.parsedate_to_datetime(value).timestamp() - time.time()
348 except (TypeError, ValueError):
349 return None
350 return min(MAX_RETRY_AFTER, max(0.0, seconds))
353def _download_once(
354 url: str,
355 dest: Path,
356 show_progress: bool = True,
357 expected_md5: str | None = None,
358 resume: bool = True,
359 timeout: float = 60.0,
360 progress: Progress | None = None,
361) -> Path:
362 """Download *url* to *dest* via ``<dest>.part`` with HTTP Range resume.
364 * If a ``.part`` file exists and *resume* is set, a ``Range: bytes=N-``
365 request continues it. A ``206`` reply is appended; a ``200`` reply
366 (server ignored the range) restarts from zero; ``416`` means the part
367 is already complete.
368 * When *expected_md5* is given the finished file is verified before the
369 atomic rename; on mismatch the part file is deleted and
370 [`ChecksumError`][lunar_toolkit.downloader.ChecksumError] is raised.
371 * On network errors the ``.part`` file is kept so the next call resumes.
372 """
373 from tqdm import tqdm
375 url = modernize_pds_url(url)
376 dest = Path(dest)
377 dest.parent.mkdir(parents=True, exist_ok=True)
378 tmp = dest.with_suffix(dest.suffix + ".part")
380 offset = tmp.stat().st_size if (resume and tmp.exists()) else 0
381 req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT})
382 if offset:
383 req.add_header("Range", f"bytes={offset}-")
385 log.info("Downloading %s%s", url, f" (resuming at {offset:,} bytes)" if offset else "")
386 try:
387 resp = urllib.request.urlopen(req, timeout=timeout)
388 except urllib.error.HTTPError as exc:
389 if exc.code != 416 or not offset:
390 raise
391 resp = None # 416: nothing left to fetch — the part file is complete
393 if resp is not None:
394 with resp:
395 status = getattr(resp, "status", 200)
396 if offset and status != 206:
397 offset = 0 # server ignored Range → start over
398 length = int(resp.headers.get("Content-Length", 0)) or None
399 total = (length + offset) if length else None
400 with (
401 open(tmp, "ab" if offset else "wb") as f,
402 tqdm(
403 total=total,
404 initial=offset,
405 unit="B",
406 unit_scale=True,
407 unit_divisor=1024,
408 desc=dest.name,
409 disable=not show_progress,
410 leave=True,
411 ) as bar,
412 ):
413 chunk = 1 << 20 # 1 MB
414 done = offset
415 if progress:
416 progress(dest.name, done, total)
417 while True:
418 buf = resp.read(chunk)
419 if not buf:
420 break
421 f.write(buf)
422 bar.update(len(buf))
423 done += len(buf)
424 if progress:
425 progress(dest.name, done, total)
426 if total is not None and tmp.stat().st_size < total: 426 ↛ 427line 426 didn't jump to line 427 because the condition on line 426 was never true
427 raise OSError(
428 f"Incomplete download of {url}: {tmp.stat().st_size:,}/{total:,} bytes "
429 f"(partial file kept at {tmp}; re-run to resume)"
430 )
432 if expected_md5:
433 actual = file_md5(tmp)
434 if actual.lower() != expected_md5.lower():
435 tmp.unlink()
436 raise ChecksumError(f"MD5 mismatch for {url}: expected {expected_md5}, got {actual}")
437 tmp.replace(dest)
438 return dest
441def _download_one(
442 url: str, dest: Path, show_progress: bool, force: bool = False
443) -> tuple[str, Path | None, str | None]:
444 if reuse_existing(dest, force):
445 return url, dest, None
446 try:
447 return url, download_file(url, dest, show_progress=show_progress), None
448 except Exception as exc:
449 return url, None, str(exc)