PVNet 4.0.9__tar.gz → 4.0.11__tar.gz

{pvnet-4.0.9 → pvnet-4.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PVNet
-Version: 4.0.9
+Version: 4.0.11
 Summary: PVNet
 Author-email: Peter Dudfield <info@openclimatefix.org>
 Description-Content-Type: text/markdown

{pvnet-4.0.9 → pvnet-4.0.11}/PVNet.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PVNet
-Version: 4.0.9
+Version: 4.0.11
 Summary: PVNet
 Author-email: Peter Dudfield <info@openclimatefix.org>
 Description-Content-Type: text/markdown

pvnet-4.0.11/pvnet/__init__.py

@@ -0,0 +1,2 @@
+"""PVNet"""
+__version__ = "4.0.11"

{pvnet-4.0.9 → pvnet-4.0.11}/pvnet/models/base_model.py

@@ -4,6 +4,7 @@ import json
 import logging
 import os
 import tempfile
+import time
 from pathlib import Path
 from typing import Dict, Optional, Union
 
@@ -141,6 +142,65 @@ def minimize_data_config(input_path, output_path, model):
         yaml.dump(config, outfile, default_flow_style=False)
 
 
+def download_hf_hub_with_retries(
+    repo_id,
+    filename,
+    revision,
+    cache_dir,
+    force_download,
+    proxies,
+    resume_download,
+    token,
+    local_files_only,
+    max_retries=5,
+    wait_time=10,
+):
+    """
+    Tries to download a file from HuggingFace up to max_retries times.
+
+    Args:
+        repo_id (str): HuggingFace repo ID
+        filename (str): Name of the file to download
+        revision (str): Specific model revision
+        cache_dir (str): Cache directory
+        force_download (bool): Whether to force a new download
+        proxies (dict): Proxy settings
+        resume_download (bool): Resume interrupted downloads
+        token (str): HuggingFace auth token
+        local_files_only (bool): Use local files only
+        max_retries (int): Maximum number of retry attempts
+        wait_time (int): Wait time (in seconds) before retrying
+
+    Returns:
+        str: The local file path of the downloaded file
+    """
+    for attempt in range(1, max_retries + 1):
+        try:
+            return hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                token=token,
+                local_files_only=local_files_only,
+            )
+        except Exception as e:
+            if attempt == max_retries:
+                raise Exception(
+                    f"Failed to download {filename} from {repo_id} after {max_retries} attempts."
+                ) from e
+            logging.warning(
+                (
+                    f"Attempt {attempt}/{max_retries} failed to download {filename} "
+                    f"from {repo_id}. Retrying in {wait_time} seconds..."
+                )
+            )
+            time.sleep(wait_time)
+
+
 class PVNetModelHubMixin(PyTorchModelHubMixin):
     """
     Implementation of [`PyTorchModelHubMixin`] to provide model Hub upload/download capabilities.
@@ -169,7 +229,7 @@ class PVNetModelHubMixin(PyTorchModelHubMixin):
             config_file = os.path.join(model_id, CONFIG_NAME)
         else:
             # load model file
-            model_file = hf_hub_download(
+            model_file = download_hf_hub_with_retries(
                 repo_id=model_id,
                 filename=PYTORCH_WEIGHTS_NAME,
                 revision=revision,
@@ -179,10 +239,12 @@ class PVNetModelHubMixin(PyTorchModelHubMixin):
                 resume_download=resume_download,
                 token=token,
                 local_files_only=local_files_only,
+                max_retries=5,
+                wait_time=10,
             )
 
             # load config file
-            config_file = hf_hub_download(
+            config_file = download_hf_hub_with_retries(
                 repo_id=model_id,
                 filename=CONFIG_NAME,
                 revision=revision,
@@ -192,6 +254,8 @@ class PVNetModelHubMixin(PyTorchModelHubMixin):
                 resume_download=resume_download,
                 token=token,
                 local_files_only=local_files_only,
+                max_retries=5,
+                wait_time=10,
             )
 
         with open(config_file, "r", encoding="utf-8") as f:
@@ -222,7 +286,7 @@ class PVNetModelHubMixin(PyTorchModelHubMixin):
             print("Loading data config from local directory")
             data_config_file = os.path.join(model_id, DATA_CONFIG_NAME)
         else:
-            data_config_file = hf_hub_download(
+            data_config_file = download_hf_hub_with_retries(
                 repo_id=model_id,
                 filename=DATA_CONFIG_NAME,
                 revision=revision,
@@ -232,6 +296,8 @@ class PVNetModelHubMixin(PyTorchModelHubMixin):
                 resume_download=resume_download,
                 token=token,
                 local_files_only=local_files_only,
+                max_retries=5,
+                wait_time=10,
             )
 
         return data_config_file

pvnet-4.0.9/pvnet/__init__.py

@@ -1,2 +0,0 @@
-"""PVNet"""
-__version__ = "4.0.9"