fix(filepath): fix the dowload as you go and delete the temps

Author: kshitijrajsharma

osmsg/pipeline.py

@@ -7,6 +7,7 @@
 import datetime as dt
 import hashlib
 import json
+import multiprocessing
 import os
 import shutil
 from dataclasses import dataclass, field
@@ -543,49 +544,81 @@ def _processing_config(cfg: RunConfig, *, parquet_dir: Path, geom_wkt: str | Non
 _DOWNLOAD_WORKERS = 4
 
 
-def _download_all(
+def _stream_window(url: str) -> int:
+    """Cap on raw files held on disk at once, sized by file weight: planet day diffs are ~1 GB so
+    only a few fit; hour diffs are moderate; minute and changeset files are tiny so many can."""
+    lowered = url.lower()
+    if "day" in lowered:
+        return 4
+    if "hour" in lowered:
+        return 24
+    return 100
+
+
+def _stream_download_process(
     urls: list[str],
+    *,
     mode: str,
-    workers: int,
     cookie: str | None,
     cache_dir: Path,
-    label: str,
-    description: str = "downloading",
-) -> None:
-    try:
-        with (
-            progress_bar(len(urls), unit=label, description=description) as advance,
-            concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool,
-        ):
-            for _ in pool.map(lambda u: download_osm_file(u, mode=mode, cookie=cookie, cache_dir=cache_dir), urls):
-                advance()
-    except requests.exceptions.RequestException as exc:
-        raise OsmsgError(
-            f"Network error downloading {label} after retries ({type(exc).__name__}). "
-            "Re-run to resume: finished downloads are cached, so it continues from where it stopped."
-        ) from exc
-
-
-def _process_all(
-    items: list,
-    *,
+    window: int,
     target,
     initializer,
-    init_args,
+    init_args: tuple,
     chunksize: int,
-    label: str,
     workers: int,
-    extra_iterables: tuple[list, ...] = (),
-    description: str = "processing",
+    label: str,
+    description: str,
+    extra_iterable: list | None = None,
 ) -> None:
+    """Download and parse the range in overlapping windows so the next window downloads while the
+    current one is parsed and its raw files are erased. Peak disk stays ~one window, not the range."""
+    total = len(urls)
+
+    def download_window(download_pool: concurrent.futures.ThreadPoolExecutor, start: int, end: int) -> None:
+        try:
+            for _ in download_pool.map(
+                lambda file_url: download_osm_file(file_url, mode=mode, cookie=cookie, cache_dir=cache_dir),
+                urls[start:end],
+            ):
+                pass
+        except requests.exceptions.RequestException as exc:
+            raise OsmsgError(
+                f"Network error downloading {label} after retries ({type(exc).__name__}). "
+                "Re-run to resume: finished downloads are cached, so it continues from where it stopped."
+            ) from exc
+
     with (
-        progress_bar(len(items), unit=label, description=description) as advance,
+        progress_bar(total, unit=label, description=description) as advance,
+        concurrent.futures.ThreadPoolExecutor(max_workers=_DOWNLOAD_WORKERS) as downloaders,
+        concurrent.futures.ThreadPoolExecutor(max_workers=1) as prefetcher,
+        # spawn, not fork: this pool lives alongside the download threads, and forking a
+        # multi-threaded process risks inheriting a held lock and deadlocking.
         concurrent.futures.ProcessPoolExecutor(
-            max_workers=workers, initializer=initializer, initargs=init_args
-        ) as pool,
+            max_workers=workers,
+            mp_context=multiprocessing.get_context("spawn"),
+            initializer=initializer,
+            initargs=init_args,
+        ) as processors,
     ):
-        for _ in pool.map(target, items, *extra_iterables, chunksize=chunksize):
-            advance()
+        download_window(downloaders, 0, min(window, total))
+        start = 0
+        while start < total:
+            end = min(start + window, total)
+            next_start = end
+            prefetching = (
+                prefetcher.submit(download_window, downloaders, next_start, min(next_start + window, total))
+                if next_start < total
+                else None
+            )
+            process_iterables = (
+                (urls[start:end],) if extra_iterable is None else (urls[start:end], extra_iterable[start:end])
+            )
+            for _ in processors.map(target, *process_iterables, chunksize=chunksize):
+                advance()
+            if prefetching is not None:
+                prefetching.result()
+            start = end
 
 
 def run(cfg: RunConfig) -> dict[str, Any]:
@@ -755,24 +788,19 @@ def run(cfg: RunConfig) -> dict[str, Any]:
             cs_config = _processing_config(cfg, parquet_dir=cs_dir, geom_wkt=geom_wkt)
             cs_config["window_start_utc"] = cfg.start_date.astimezone(UTC)
 
-            _download_all(
-                urls,
-                "changeset",
-                _DOWNLOAD_WORKERS,
-                None,
-                cfg.cache_dir,
-                "changesets",
-                description="Downloading changesets",
-            )
-            _process_all(
+            _stream_download_process(
                 urls,
+                mode="changeset",
+                cookie=None,
+                cache_dir=cfg.cache_dir,
+                window=_stream_window(CHANGESETS_REPLICATION),
                 target=process_changeset,
                 initializer=init_changeset_worker,
                 init_args=(cs_config,),
                 chunksize=10,
-                label="changesets",
                 workers=max_workers,
-                description="Processing changesets",
+                label="changesets",
+                description="Changesets",
             )
             dbmod.merge_parquet_files(conn, cs_dir, cleanup=True)
             upsert_state(
@@ -822,27 +850,22 @@ def run(cfg: RunConfig) -> dict[str, Any]:
         cf_config = _processing_config(cfg, parquet_dir=cf_dir, geom_wkt=None)
         cf_config["start_date_utc"] = url_start_date_utc
 
-        _download_all(
-            urls,
-            "changefiles",
-            _DOWNLOAD_WORKERS,
-            cookie,
-            cfg.cache_dir,
-            "changefiles",
-            description="Downloading changefiles",
-        )
         chunksize = 10 if "minute" in url.lower() else 1
         seq_ids = list(range(src_start_seq, src_end_seq + 1))
-        _process_all(
+        _stream_download_process(
             urls,
+            mode="changefiles",
+            cookie=cookie,
+            cache_dir=cfg.cache_dir,
+            window=_stream_window(url),
             target=process_changefile,
             initializer=init_changefile_worker,
             init_args=(valid_changesets, cf_config),
             chunksize=chunksize,
-            label="changefiles",
             workers=max_workers,
-            extra_iterables=(seq_ids,),
-            description="Processing changefiles",
+            label="changefiles",
+            description="Changefiles",
+            extra_iterable=seq_ids,
         )
         dbmod.merge_parquet_files(conn, cf_dir, cleanup=True)
         # state.last_ts is the seq_ts of last_seq so the next tick's lower-bound filter

tests/test_download_resilience.py

@@ -11,17 +11,23 @@ def test_download_concurrency_is_capped():
     assert pipeline._DOWNLOAD_WORKERS <= 4
 
 
-def test_download_all_wraps_network_error(monkeypatch, tmp_path):
+def test_stream_download_wraps_network_error(monkeypatch, tmp_path):
     def boom(url, **kwargs):
         raise requests.exceptions.ConnectTimeout("planet timed out")
 
     monkeypatch.setattr(pipeline, "download_osm_file", boom)
     with pytest.raises(OsmsgError, match="Re-run to resume"):
-        pipeline._download_all(
+        pipeline._stream_download_process(
             ["https://planet.openstreetmap.org/replication/changesets/007/035/882.osm.gz"],
-            "changeset",
-            8,
-            None,
-            tmp_path,
-            "changesets",
+            mode="changeset",
+            cookie=None,
+            cache_dir=tmp_path,
+            window=250,
+            target=int,
+            initializer=None,
+            init_args=(),
+            chunksize=1,
+            workers=1,
+            label="changesets",
+            description="Changesets",
         )

tests/test_pipeline_smoke.py

@@ -102,8 +102,7 @@ def fake_changefile_download_urls(_start, _end, _base, *, resume_seq=None, cs_ts
         return (["https://fake/101.osc.gz"], fake_server_ts, 101, 101, "u1", "u2")
 
     monkeypatch.setattr(pipeline_mod, "changefile_download_urls", fake_changefile_download_urls)
-    monkeypatch.setattr(pipeline_mod, "_download_all", lambda *a, **kw: None)
-    monkeypatch.setattr(pipeline_mod, "_process_all", lambda *a, **kw: None)
+    monkeypatch.setattr(pipeline_mod, "_stream_download_process", lambda *a, **kw: None)
     monkeypatch.setattr(pipeline_mod.dbmod, "merge_parquet_files", lambda *a, **kw: None)
     monkeypatch.setattr(pipeline_mod, "changefile_seq_timestamp", lambda _base, _seq: seq_ts)
     monkeypatch.setattr(

tests/test_pipeline_stream.py

@@ -0,0 +1,112 @@
+"""_stream_download_process streams in overlapping windows: full coverage, seq alignment,
+download-before-process ordering, and a bounded on-disk footprint (delete-as-you-go)."""
+
+import pathlib
+
+import pytest
+import requests
+
+from osmsg import pipeline
+from osmsg.exceptions import OsmsgError
+
+_RESULTS = ""
+_CACHE = ""
+
+
+def _init(results_dir: str, cache_dir: str) -> None:
+    global _RESULTS, _CACHE
+    _RESULTS, _CACHE = results_dir, cache_dir
+
+
+def _probe(url: str, seq_id: int | None = None) -> None:
+    base = url.rsplit("/", 1)[-1]
+    raw = pathlib.Path(_CACHE) / base
+    existed = raw.exists()
+    on_disk = sum(1 for f in pathlib.Path(_CACHE).iterdir() if f.is_file())
+    key = seq_id if seq_id is not None else base
+    (pathlib.Path(_RESULTS) / f"r_{key}").write_text(f"{url}|existed={existed}|on_disk={on_disk}")
+    if existed:
+        raw.unlink()  # mimic the worker's per-file erase
+
+
+def _fake_download(url: str, mode: str = "", cookie: str | None = None, cache_dir=None) -> pathlib.Path:
+    path = pathlib.Path(cache_dir) / url.rsplit("/", 1)[-1]
+    path.write_text("x")
+    return path
+
+
+def _run(urls, extra, window, workers, cache, results) -> None:
+    pipeline._stream_download_process(
+        urls,
+        mode="changefiles",
+        cookie=None,
+        cache_dir=pathlib.Path(cache),
+        window=window,
+        target=_probe,
+        initializer=_init,
+        init_args=(str(results), str(cache)),
+        chunksize=1,
+        workers=workers,
+        label="files",
+        description="files",
+        extra_iterable=extra,
+    )
+
+
+def test_streams_covered_aligned_and_bounded(monkeypatch, tmp_path):
+    monkeypatch.setattr(pipeline, "download_osm_file", _fake_download)
+    cache, results = tmp_path / "cache", tmp_path / "results"
+    cache.mkdir()
+    results.mkdir()
+
+    count, window = 23, 4
+    urls = [f"https://x/{i:03d}.osc.gz" for i in range(count)]
+    seqs = list(range(1000, 1000 + count))
+    _run(urls, seqs, window, 3, cache, results)
+
+    peak = 0
+    for url, seq in zip(urls, seqs):
+        body = (results / f"r_{seq}").read_text()
+        assert body.startswith(url + "|"), f"seq {seq} misaligned: {body}"
+        assert "existed=True" in body, f"processed before download: {body}"
+        peak = max(peak, int(body.split("on_disk=")[1]))
+
+    assert len(list(results.iterdir())) == count, "every file must be processed exactly once"
+    assert peak <= 2 * window, f"footprint {peak} exceeds one prefetch window ahead ({2 * window})"
+    assert not list(cache.iterdir()), "raw files must be erased as they are processed"
+
+
+def test_streams_without_extra_iterable(monkeypatch, tmp_path):
+    monkeypatch.setattr(pipeline, "download_osm_file", _fake_download)
+    cache, results = tmp_path / "cache", tmp_path / "results"
+    cache.mkdir()
+    results.mkdir()
+    urls = [f"https://x/{i:03d}.osm.gz" for i in range(15)]
+    _run(urls, None, 100, 2, cache, results)
+    assert len(list(results.iterdir())) == 15
+
+
+def test_single_window(monkeypatch, tmp_path):
+    monkeypatch.setattr(pipeline, "download_osm_file", _fake_download)
+    cache, results = tmp_path / "cache", tmp_path / "results"
+    cache.mkdir()
+    results.mkdir()
+    urls = [f"https://x/{i:03d}.osc.gz" for i in range(3)]
+    _run(urls, [0, 1, 2], 8, 2, cache, results)
+    assert len(list(results.iterdir())) == 3
+
+
+def test_network_error_is_wrapped(monkeypatch, tmp_path):
+    def boom(url, **kwargs):
+        raise requests.exceptions.ConnectTimeout("down")
+
+    monkeypatch.setattr(pipeline, "download_osm_file", boom)
+    with pytest.raises(OsmsgError, match="Re-run to resume"):
+        _run(["https://x/1.osc.gz"], [1], 4, 1, tmp_path / "cache", tmp_path)
+
+
+def test_stream_window_sizes_by_file_weight():
+    assert pipeline._stream_window("https://planet/replication/day") == 4
+    assert pipeline._stream_window("https://planet/replication/hour") == 24
+    assert pipeline._stream_window("https://planet/replication/minute") == 100
+    assert pipeline._stream_window("https://planet/replication/changesets/") == 100