Don't require tqdm for downloads

valohai/internals/download.py

@@ -1,3 +1,4 @@
+import contextlib
 import os
 import tempfile
 
@@ -15,25 +16,25 @@ def download_url(url: str, path: str, force_download: bool = False) -> str:
 def _do_download(url: str, path: str) -> None:
     try:
         import requests
-        from tqdm import tqdm
     except ImportError as ie:
         raise RuntimeError(
-            f"The `requests` and `tqdm` modules must be available "
-            f"for download support (attempting to download {url})"
+            f"The `requests` module must be available for download support (attempting to download {url})"
         ) from ie
 
     tmp_path = tempfile.NamedTemporaryFile().name
     print(f"Downloading {url} -> {path}")  # noqa
     r = requests.get(url, stream=True)
     r.raise_for_status()
-    total = (
-        int(r.headers["content-length"]) if "content-length" in r.headers else None
-    )
+    total = int(r.headers["content-length"]) if "content-length" in r.headers else None
+    try:
+        from tqdm import tqdm
+        prog = tqdm(total=total, unit="iB", unit_scale=True)
+    except ImportError:
+        prog = contextlib.nullcontext()
 
-    with tqdm(total=total, unit="iB", unit_scale=True) as t, open(
-        tmp_path, "wb"
-    ) as f:
-        for data in r.iter_content(1048576):
-            t.update(len(data))
-            f.write(data)
+    with prog, open(tmp_path, "wb") as f:
+        for chunk in r.iter_content(1048576):
+            if prog:
+                prog.update(len(chunk))
+            f.write(chunk)
     os.replace(tmp_path, path)