Massively sped up puller by threading links lookup, and improved error handling by verifying file sizes.

puller_fits.py

@@ -8,22 +8,36 @@ import datetime
 import json
 from threading import Thread
 import queue
+import math
+from functools import partial
 
 import tqdm
 
 directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/"
 stored_images_dir = r"..\Data"
-# directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/goes16/l2/data/suvi-l2-ci094/"
+# directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/goes16/l2/data/ephe-l2-orb1m/2017/12/"
-# stored_images_dir = r"Z:\NOAA GOES Data\Data\goes16\l2\suvi-l2-ci094"
+# stored_images_dir = r"Z:\NOAA GOES Data\Data\goes16/l2/data/ephe-l2-orb1m/2017/12/"
 ignore_folder_names = ["l1b", "goes17", "2017", "2018", "2019", "2020", "2021", "2022", "2023"]
 file_database_path = r"..\file_database.json"
 fetch_interval = 0 # 60*60
-nworkers = 8
+nfetchworkers = 10 # Be nice to the servers, this value is how many threads will be asking for links and file info at the same time
+ndownloadworkers = 4 # Be nice to the servers, this value is how many threads will be downloading files at the same time
 randomize_order = False
 
-def recursive_find_links(url, attempt_count = 3, randomize_order = True):
+links_regex_pattern = r'(?<=<a href=")([^ ?:]*)(?=">.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace
-    links_regex_pattern = r'(?<=<a href=")([^ ?:]*)(?=">.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace
+times_regex_pattern = r'(?<=<td align="right">)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm
-    times_regex_pattern = r'(?<=<td align="right">)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm
+sizes_regex_pattern = r'(?<=\d{4}-\d{2}-\d{2} \d{2}:\d{2}  <\/td><td align="right">)([ .\d]{1,3})([KMG]|- |0 | )(?=<\/td>)' # Find file size markers in one of many possible formats
+links_matcher = re.compile(links_regex_pattern)
+times_matcher = re.compile(times_regex_pattern)
+sizes_matcher = re.compile(sizes_regex_pattern)
+
+def find_links_worker(query_work_queue, query_result_queue, attempt_count = 3, randomize_order = True):
+    while True:
+        job = query_work_queue.get()
+        if job == None:
+            return
+        url= job
+    
         attempts = 0
         html_content = None
         while attempts < attempt_count:
@@ -39,50 +53,67 @@ def recursive_find_links(url, attempt_count = 3, randomize_order = True):
                 print(f'\nAfter {attempt_count} retries, could not fetch: {url}')
                 return
 
-    links = re.findall(links_regex_pattern, html_content)
+        links = links_matcher.findall(html_content)
-    times = re.findall(times_regex_pattern, html_content)
+        times = times_matcher.findall(html_content)
+        sizes = sizes_matcher.findall(html_content)
         for i in range(len(times)):
             dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M")
             dt.replace(tzinfo=datetime.timezone.utc)
             times[i] = time.mktime(dt.timetuple())
+        for i in range(len(sizes)):
+            match sizes[i][1].strip():
+                case '-':
+                    sizes[i] = 0
+                case 'K':
+                    sizes[i] = int(float(sizes[i][0].strip()) * 1024)
+                case 'M':
+                    sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024)
+                case 'G':
+                    sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024 * 1024)
+                case "0":
+                    sizes[i] = 0
+                case "":
+                    sizes[i] = int(float(sizes[i][0].strip()))
+                case _:
+                    raise(ValueError(f"Unexpected symbol while parsing links page: {_}"))
 
-    if len(links) - 1 == len(times):
-        links = links[1:]
-    elif len(links) == len(times):
-        pass
-    else:
-        raise(ValueError)
         
-    results = list(zip(links, times))
+        if (len(links) != len(times)) or (len(times) != len(sizes)):
+            raise(ValueError("Links parsing error, mismatched numbers of links, times, or sizes!"))
+
+        results = list(zip(links, times, sizes))
         if randomize_order:
             random.shuffle(results)
-    for _link, _time in results:
+        for _link, _time, _size in results:
             if _link.endswith("/"):
                 if _link.split(r"/")[-2] in ignore_folder_names:
                     continue
                 else:
-                yield from recursive_find_links(urllib.parse.urljoin(url,_link), attempt_count, randomize_order)
+                    query_work_queue.put(urllib.parse.urljoin(url,_link))
             else:
-            yield urllib.parse.urljoin(url,_link), _time
+                query_result_queue.put((urllib.parse.urljoin(url,_link), _time, _size))
     
 
-# Fetch image from url and store it to path, retrying on failure
+# Fetch file from url and store it to path, retrying on failure
-def image_fetch_worker(work_queue, result_queue, attempt_count = 1):
+def file_download_worker(download_work_queue, download_result_queue, attempt_count = 2):
     while True:
-        job = work_queue.get()
+        job = download_work_queue.get()
         if job == None:
             return
-        url, path, t = job
+        url, path, t, s = job
 
         attempts = 0
         while attempts < attempt_count:
             try:
                 req = urllib.request.Request(url, data=None)
                 image_data = urllib.request.urlopen(req).read()
+                if math.isclose(len(image_data), s, rel_tol=0.05):
                     os.makedirs(os.path.split(path)[0], exist_ok=True)
                     open(path, 'wb').write(image_data)
-                result_queue.put((True, url, t))
+                    download_result_queue.put((True, url, t))
                     break
+                else:
+                    raise ValueError("Downloaded file is the wrong size!")
             except Exception as e:
                 if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov)
                     attempts += attempt_count
@@ -92,7 +123,7 @@ def image_fetch_worker(work_queue, result_queue, attempt_count = 1):
                     attempts += 1
             if (attempt_count > 1) and (attempts == attempt_count):
                 print(f'\nAfter {attempt_count} retries, could not fetch: {url}')
-                result_queue.put((False, url, t))
+                download_result_queue.put((False, url, t))
                 break
 
 
@@ -107,24 +138,33 @@ if __name__ == "__main__":
         print(f"Load failed, starting with empty cache")
         file_info_cache = {}
 
-    work_queue = queue.Queue(maxsize=nworkers)
+    query_work_queue = queue.Queue()
-    result_queue = queue.Queue()
+    query_result_queue = queue.Queue()
+    download_work_queue = queue.Queue()
+    download_result_queue = queue.Queue()
     workers = []
-    for _ in range(nworkers):
+    for _ in range(nfetchworkers):
-        t = Thread(target=image_fetch_worker, args=(work_queue, result_queue), daemon=True)
+        t = Thread(target=find_links_worker, args=(query_work_queue, query_result_queue, 3, randomize_order), daemon=True)
         t.start()
         workers.append(t)
+    for _ in range(ndownloadworkers):
+        t = Thread(target=file_download_worker, args=(download_work_queue, download_result_queue), daemon=True)
+        t.start()
+        workers.append(t)
+    
     try:
         while True:
             fetched_image_count = 0
             already_had_image_count = 0
             failed_image_count = 0
             urllen = len(directory_url)
-            for l, t in tqdm.tqdm(recursive_find_links(directory_url, randomize_order = randomize_order), desc="Downloading files"):
+
+            query_work_queue.put(directory_url)
+            for l, t, s in tqdm.tqdm(iter(partial(query_result_queue.get, timeout=30.0), None), desc="Downloading files"):
                 # Collect completed jobs and record completion status
                 while True:
                     try:
-                        r_success, r_url, r_t = result_queue.get_nowait()
+                        r_success, r_url, r_t = download_result_queue.get_nowait()
                         if r_success: 
                             fetched_image_count += 1
                             file_info_cache[r_url] = r_t
@@ -155,11 +195,15 @@ if __name__ == "__main__":
                         if l in file_info_cache: # If we have record of this file, it must be out of date, delete it and download the new version.
                             if os.path.exists(filepath):
                                 os.remove(filepath)
-                        else: # If we have no record of this file, update the file info cache and don't redownload
+                        else: # If we have no record of this file, update the file info cache and don't redownload if fsize is right
+                            fsize = os.path.getsize(filepath)
+                            if math.isclose(fsize, s, rel_tol=0.05):
                                 file_info_cache[l] = t
                                 already_had_image_count += 1
                                 continue
-                    work_queue.put((l, filepath, t))
+                            else:
+                                print(f'Found a mismatched size on file: {filepath}  Redownloading!')
+                    download_work_queue.put((l, filepath, t, s))
                 else:
                     # We have a download record, confirm the file actually exists on disk
                     if os.path.exists(filepath):
@@ -170,7 +214,7 @@ if __name__ == "__main__":
                         if os.path.exists(filepath2) or os.path.exists(filepath3):
                             already_had_image_count += 1
                         else: # We could not find the file on disk, queue for redownload
-                            work_queue.put((l, filepath, t))
+                            download_work_queue.put((l, filepath, t))
                     
             with open(file_database_path, 'w') as f:
                 f.write(json.dumps(file_info_cache))
@@ -185,10 +229,19 @@ if __name__ == "__main__":
                 break
     except KeyboardInterrupt:
         print("Saving file database and shutting down.")
+    except Exception as e:
+        print(f"Unhandled Exception during run: {e}")
+        print("Shutting down")
 
-    for _ in range(nworkers):
+    for _ in range(nfetchworkers):
         try:
-            work_queue.put(None, timeout=5.0)
+            query_work_queue.put(None, timeout=5.0)
+        except:
+            break
+
+    for _ in range(ndownloadworkers):
+        try:
+            download_work_queue.put(None, timeout=5.0)
         except:
             break