Retry on transient error codes in preload (#5982)

distributed/preloading.py

@@ -6,13 +6,13 @@
 import os
 import shutil
 import sys
-import urllib.request
 from collections.abc import Iterable
 from importlib import import_module
 from types import ModuleType
 from typing import TYPE_CHECKING, cast
 
 import click
+import urllib3
 
 from dask.utils import tmpfile
 
@@ -131,9 +131,17 @@ def _download_module(url: str) -> ModuleType:
     logger.info("Downloading preload at %s", url)
     assert is_webaddress(url)
 
-    request = urllib.request.Request(url, method="GET")
-    response = urllib.request.urlopen(request)
-    source = response.read().decode()
+    with urllib3.PoolManager() as http:
+        response = http.request(
+            method="GET",
+            url=url,
+            retries=urllib3.util.Retry(
+                status_forcelist=[429, 504, 503, 502],
+                backoff_factor=0.2,
+            ),
+        )
+
+        source = response.data
 
     compiled = compile(source, url, "exec")
     module = ModuleType(url)

requirements.txt

@@ -9,5 +9,6 @@ sortedcontainers !=2.0.0, !=2.0.1
 tblib >= 1.6.0
 toolz >= 0.8.2
 tornado >= 6.0.3
+urllib3
 zict >= 0.1.3
 pyyaml