duplicate download util for shared template

truss/templates/shared/lazy_data_resolver.py

@@ -4,7 +4,7 @@
 
 import pydantic
 import yaml
-from truss.util.download import download_from_url_using_requests
+from shared.util import download_from_url_using_requests
 
 LAZY_DATA_RESOLVER_PATH = Path("/bptr/bptr-manifest")
 

truss/templates/shared/util.py

@@ -1,10 +1,14 @@
 import multiprocessing
 import os
+import shutil
 import sys
+from pathlib import Path
 from typing import Callable, Dict, List, TypeVar
 
 import psutil
+import requests
 
+BLOB_DOWNLOAD_TIMEOUT_SECS = 600  # 10 minutes
 # number of seconds to wait for truss server child processes before sending kill signal
 CHILD_PROCESS_WAIT_TIMEOUT_SECONDS = 120
 
@@ -85,3 +89,16 @@ def kill_child_processes(parent_pid: int):
 
 def transform_keys(d: Dict[X, Z], fn: Callable[[X], Y]) -> Dict[Y, Z]:
     return {fn(key): value for key, value in d.items()}
+
+
+def download_from_url_using_requests(URL: str, download_to: Path):
+    # Streaming download to keep memory usage low
+    resp = requests.get(
+        URL,
+        allow_redirects=True,
+        stream=True,
+        timeout=BLOB_DOWNLOAD_TIMEOUT_SECS,
+    )
+    resp.raise_for_status()
+    with download_to.open("wb") as file:
+        shutil.copyfileobj(resp.raw, file)