implement resuming ahn downloads on retry

Author: Ylannl

packages/common/src/bag3d/common/utils/requests.py

@@ -3,6 +3,7 @@
 from time import sleep
 from typing import Mapping, Union
 from urllib.parse import urlparse, urljoin
+import os
 
 import requests
 from dagster import (
@@ -35,6 +36,7 @@ def download_file(
     chunk_size: int = 1024,
     parameters: Mapping = None,
     verify: bool = True,
+    attempt_resume: bool = False,
 ) -> Union[Path, None]:
     """Download a large file and save it to disk.
 
@@ -59,6 +61,42 @@ def download_file(
     session = requests.Session()  # https://stackoverflow.com/a/63417213
 
     try:
+        head = session.head(url, allow_redirects=True)
+        head.raise_for_status()
+        remote_size = int(head.headers.get("content-length", 0))
+        headers = {}
+        if attempt_resume and os.path.exists(fpath):
+            local_size = os.path.getsize(fpath)
+            headers = {"Range": f"bytes={local_size}-"}
+            logger.info(f"Resuming download at {local_size}/{remote_size} bytes")
+        else:
+            local_size = 0
+            logger.info(f"Starting download of {remote_size} bytes")
+
+        with fpath.open("ab") as fd:
+            with session.get(
+                url, headers=headers, stream=True, verify=verify, params=parameters
+            ) as r:
+                if local_size and r.status_code != 206:
+                    raise ValueError("Server does not support range requests")
+                elif r.status_code not in (200, 206):
+                    r.raise_for_status()
+
+                bytes_written = local_size
+                last_logged_percent = (
+                    int((bytes_written / remote_size) * 100) if remote_size else 0
+                )
+                for data in r.iter_content(chunk_size=chunk_size):
+                    fd.write(data)
+                    bytes_written += len(data)
+                    if remote_size:
+                        percent = int((bytes_written / remote_size) * 100)
+                        if percent > last_logged_percent and percent % 10 == 0:
+                            logger.info(
+                                f"Download progress: {percent}% ({bytes_written}/{remote_size} bytes)"
+                            )
+                            last_logged_percent = percent
+
         r = session.get(url, params=parameters, stream=True, verify=verify)
         if r.ok:
             with fpath.open("wb") as fd:
@@ -69,9 +107,11 @@ def download_file(
             r.raise_for_status()
         r.close()
     except (
+        requests.RequestException,
         requests.exceptions.BaseHTTPError,
         requests.exceptions.HTTPError,
         requests.exceptions.ChunkedEncodingError,
+        ValueError,
     ) as e:  # pragma: no cover
         logger.exception(e)
         return None

packages/core/src/bag3d/core/assets/ahn/download.py

@@ -158,7 +158,7 @@ class LazFilesConfig(Config):
     required_resource_keys={"file_store"},
     partitions_def=partition_definition_ahn,
     tags={"dagster/concurrency_key": "laz_download"},
-    pool="laz_download"
+    pool="laz_download",
 )
 def laz_files_ahn3(context, config: LazFilesConfig, md5_ahn3, tile_index_ahn):
     """AHN3 LAZ files as they are downloaded from PDOK.
@@ -213,7 +213,7 @@ def laz_files_ahn3(context, config: LazFilesConfig, md5_ahn3, tile_index_ahn):
     required_resource_keys={"file_store"},
     partitions_def=partition_definition_ahn,
     tags={"dagster/concurrency_key": "laz_download"},
-    pool="laz_download"
+    pool="laz_download",
 )
 def laz_files_ahn4(context, config: LazFilesConfig, md5_ahn4, tile_index_ahn):
     """AHN4 LAZ files as they are downloaded from PDOK.
@@ -271,7 +271,7 @@ def laz_files_ahn4(context, config: LazFilesConfig, md5_ahn4, tile_index_ahn):
     required_resource_keys={"file_store"},
     partitions_def=partition_definition_ahn,
     tags={"dagster/concurrency_key": "laz_download"},
-    pool="laz_download"
+    pool="laz_download",
 )
 def laz_files_ahn5(context, config: LazFilesConfig, sha256_ahn5, tile_index_ahn):
     """AHN5 LAZ files as they are downloaded from PDOK.
@@ -421,6 +421,7 @@ def download_laz(
             target_path=fpath,
             chunk_size=1024 * 1024,
             verify=verify_ssl,
+            attempt_resume=True,
         )
         if fpath_download is None:
             # Download failed