diff --git a/src/mccode_antlr/cli/cache.py b/src/mccode_antlr/cli/cache.py index a0a3c46f..10ab87cf 100644 --- a/src/mccode_antlr/cli/cache.py +++ b/src/mccode_antlr/cli/cache.py @@ -226,6 +226,7 @@ def warm_registries_via_pooch(registries) -> tuple[int, int]: ------- (total_fetched, error_count) """ + from mccode_antlr.reader.registry import pooch_fetch total = errors = 0 seen: set = set() @@ -243,7 +244,7 @@ def warm_registries_via_pooch(registries) -> tuple[int, int]: print(f" [{reg.name}] downloading {len(files)} files …", flush=True) for fname in files: try: - p.fetch(fname) + pooch_fetch(p, fname) total += 1 except Exception as exc: print(f" WARNING: could not fetch {fname}: {exc}", flush=True) diff --git a/src/mccode_antlr/reader/registry.py b/src/mccode_antlr/reader/registry.py index 316feb16..af83b0bf 100644 --- a/src/mccode_antlr/reader/registry.py +++ b/src/mccode_antlr/reader/registry.py @@ -61,6 +61,15 @@ def _dedupe_identical_paths(paths: list[Path]) -> Path | None: return sorted(paths, key=str)[0] +def pooch_fetch(pooch_instance, filename: str) -> Path: + """pooch_instance.fetch(filename), safe for parallel downloads into one cache + (workaround until https://github.com/fatiando/pooch/pull/554 is released).""" + directory = getattr(pooch_instance, 'abspath', None) + if directory is not None: + Path(directory, filename).parent.mkdir(parents=True, exist_ok=True) + return Path(pooch_instance.fetch(filename)) + + def _dedupe_identical_registry_entries(pooch_instance, names: list[str]) -> str | None: """Like _dedupe_identical_paths, but for a pooch-backed remote registry: pooch's registry index already stores a content hash per file (`pooch_instance.registry` @@ -342,7 +351,7 @@ def is_available(self, name: str, ext: str = None, exact: bool = True): return self.pooch.registry_files(self.fullname(name, ext, exact)) def path(self, name: str, ext: str = None, exact: bool = True) -> Path: - return Path(self.pooch.fetch(self.fullname(name, ext, exact))) + return pooch_fetch(self.pooch, self.fullname(name, ext, exact)) def filenames(self) -> list[str]: return [x for x in self.pooch.registry_files] diff --git a/tests/test_pooch_fetch_parallel.py b/tests/test_pooch_fetch_parallel.py new file mode 100644 index 00000000..5304283c --- /dev/null +++ b/tests/test_pooch_fetch_parallel.py @@ -0,0 +1,50 @@ +"""Parallel downloads into one pooch cache (https://github.com/fatiando/pooch/issues/555).""" +import functools +import http.server +import multiprocessing +import threading + +import pooch + +from mccode_antlr.reader.registry import pooch_fetch + +NAMES = [f'd{i}/e{i}/f.txt' for i in range(100)] + + +def _fetch_all(args): + port, cache = args + p = pooch.create(path=cache, base_url=f'http://127.0.0.1:{port}/', + registry={name: None for name in NAMES}) + failed = [] + for name in NAMES: + try: + pooch_fetch(p, name) + except Exception as error: + failed.append(f'{name}: {error!r}') + return failed + + +class _QuietHandler(http.server.SimpleHTTPRequestHandler): + def log_message(self, *args): + pass + + +def test_parallel_fetch_into_new_directories(tmp_path): + served = tmp_path / 'served' + for name in NAMES: + (served / name).parent.mkdir(parents=True) + (served / name).write_text(name) + server = http.server.ThreadingHTTPServer( + ('127.0.0.1', 0), functools.partial(_QuietHandler, directory=str(served))) + threading.Thread(target=server.serve_forever, daemon=True).start() + level = pooch.get_logger().level + pooch.get_logger().setLevel('WARNING') + try: + with multiprocessing.Pool(8) as pool: + results = pool.map(_fetch_all, [(server.server_address[1], str(tmp_path / 'cache'))] * 8) + finally: + pooch.get_logger().setLevel(level) + server.shutdown() + assert [f for failed in results for f in failed] == [] + for name in NAMES: + assert (tmp_path / 'cache' / name).read_text() == name