Add file integrity checking for TimesFM model downloads

etna/core/utils.py

@@ -1,11 +1,17 @@
+import hashlib
 import inspect
 import json
+import os
 import pathlib
+import warnings
 import zipfile
 from copy import deepcopy
 from functools import wraps
 from typing import Any
 from typing import Callable
+from typing import Dict
+from typing import Optional
+from urllib import request
 
 from hydra_slayer import get_factory
 
@@ -72,6 +78,116 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+# Known model hashes for integrity verification
+# To add a hash for a model URL, download the file and compute its MD5 hash
+KNOWN_MODEL_HASHES: Dict[str, str] = {
+    # Add known model URL -> hash mappings here
+    # Example: "http://example.com/model.ckpt": "abcd1234...",
+}
+
+
+def verify_file_hash(file_path: str, expected_hash: Optional[str] = None) -> bool:
+    """
+    Verify file integrity using MD5 hash.
+
+    Parameters
+    ----------
+    file_path:
+        Path to the file to verify
+    expected_hash:
+        Expected MD5 hash. If None, verification is skipped.
+
+    Returns
+    -------
+    :
+        True if hash matches or no expected hash provided, False otherwise
+    """
+    if expected_hash is None:
+        return True
+
+    if not os.path.exists(file_path):
+        return False
+
+    try:
+        hash_md5 = hashlib.md5()
+        with open(file_path, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hash_md5.update(chunk)
+        file_hash = hash_md5.hexdigest()
+        return file_hash == expected_hash
+    except Exception:
+        return False
+
+
+def download_with_integrity_check(
+    url: str, 
+    destination_path: str, 
+    expected_hash: Optional[str] = None,
+    force_redownload: bool = False
+) -> None:
+    """
+    Download a file with integrity verification.
+
+    Parameters
+    ----------
+    url:
+        URL to download from
+    destination_path:
+        Local path to save the file
+    expected_hash:
+        Expected MD5 hash for verification. If None, no verification is performed.
+    force_redownload:
+        If True, download even if file exists and passes verification
+
+    Raises
+    ------
+    RuntimeError:
+        If download fails integrity check
+    """
+    # Check if file exists and verify integrity
+    if os.path.exists(destination_path) and not force_redownload:
+        if verify_file_hash(destination_path, expected_hash):
+            return  # File exists and is valid
+        else:
+            # File exists but hash doesn't match, re-download
+            if expected_hash is not None:
+                warnings.warn(
+                    f"Local file hash does not match expected hash. "
+                    f"This may indicate a corrupted download. Re-downloading from {url}"
+                )
+            os.remove(destination_path)
+
+    # Download the file
+    os.makedirs(os.path.dirname(destination_path), exist_ok=True)
+    request.urlretrieve(url=url, filename=destination_path)
+
+    # Verify the downloaded file
+    if not verify_file_hash(destination_path, expected_hash):
+        if expected_hash is not None:
+            os.remove(destination_path)
+            raise RuntimeError(
+                f"Downloaded file from {url} failed integrity check. "
+                f"This may indicate a network issue or corrupted download."
+            )
+
+
+def get_known_hash(url: str) -> Optional[str]:
+    """
+    Get known hash for a URL from the registry.
+
+    Parameters
+    ----------
+    url:
+        URL to look up
+
+    Returns
+    -------
+    :
+        Known hash for the URL, or None if not found
+    """
+    return KNOWN_MODEL_HASHES.get(url)
+
+
 def create_type_with_init_collector(type_: type) -> type:
     """Create type with init decorated with init_collector."""
     previous_frame = inspect.stack()[1]

etna/experimental/classification/predictability.py

@@ -1,10 +1,14 @@
+import os
+import warnings
 from typing import Dict
 from typing import List
+from typing import Optional
 from urllib import request
 
 import numpy as np
 from sklearn.base import ClassifierMixin
 
+from etna.core.utils import get_known_hash, verify_file_hash
 from etna.datasets import TSDataset
 from etna.experimental.classification.classification import TimeSeriesBinaryClassifier
 from etna.experimental.classification.feature_extraction.base import BaseTimeSeriesFeatureExtractor
@@ -98,7 +102,34 @@ def download_model(model_name: str, dataset_freq: str, path: str):
             If the model does not exist in s3.
         """
         url = f"http://etna-github-prod.cdn-tinkoff.ru/series_classification/22_11_2022/{dataset_freq}/{model_name}.pickle"
+        expected_hash = get_known_hash(url)
+
+        # Check if file exists and verify integrity
+        if os.path.exists(path):
+            if verify_file_hash(path, expected_hash):
+                return  # File exists and is valid (or no hash to check)
+            else:
+                # File exists but hash doesn't match, re-download
+                if expected_hash is not None:
+                    warnings.warn(
+                        f"Local model file hash does not match expected hash. "
+                        f"This may indicate a corrupted download. Re-downloading {model_name} from {url}"
+                    )
+                os.remove(path)
+
+        # Download the file
         try:
             request.urlretrieve(url=url, filename=path)
-        except Exception:
+
+            # Verify the downloaded file
+            if not verify_file_hash(path, expected_hash):
+                if expected_hash is not None:
+                    os.remove(path)
+                    raise RuntimeError(
+                        f"Downloaded model file {model_name} from {url} failed integrity check. "
+                        f"This may indicate a network issue or corrupted download."
+                    )
+        except Exception as e:
+            if expected_hash is not None and "integrity check" in str(e):
+                raise  # Re-raise integrity check errors
             raise ValueError("Model not found! Check the list of available models!")

etna/models/nn/chronos/base.py

@@ -5,13 +5,15 @@
 from pathlib import Path
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import Sequence
 from typing import Union
 from urllib import request
 
 import pandas as pd
 
 from etna import SETTINGS
+from etna.core.utils import get_known_hash, verify_file_hash
 from etna.datasets import TSDataset
 from etna.distributions import BaseDistribution
 from etna.models.base import PredictionIntervalContextRequiredAbstractModel
@@ -84,18 +86,55 @@ def _is_url(self):
         return self.path_or_url.startswith("https://") or self.path_or_url.startswith("http://")
 
     def _download_model_from_url(self) -> str:
-        """Download model from url to local cache_dir."""
+        """Download model from url to local cache_dir with integrity verification."""
         model_file = self.path_or_url.split("/")[-1]
         model_dir = model_file.split(".zip")[0]
         full_model_path = f"{self.cache_dir}/{model_dir}"
-        if not os.path.exists(full_model_path):
-            try:
-                request.urlretrieve(url=self.path_or_url, filename=model_file)
-
-                with zipfile.ZipFile(model_file, "r") as zip_ref:
-                    zip_ref.extractall(self.cache_dir)
-            finally:
-                os.remove(model_file)
+        zip_file_path = f"{self.cache_dir}/{model_file}"
+        expected_hash = get_known_hash(self.path_or_url)
+
+        # Check if extracted model directory exists and verify ZIP file integrity if it still exists
+        if os.path.exists(full_model_path):
+            if os.path.exists(zip_file_path):
+                if verify_file_hash(zip_file_path, expected_hash):
+                    return full_model_path
+                else:
+                    # ZIP file exists but hash doesn't match, re-download
+                    if expected_hash is not None:
+                        warnings.warn(
+                            f"Local model ZIP file hash does not match expected hash. "
+                            f"This may indicate a corrupted download. Re-downloading from {self.path_or_url}"
+                        )
+                    # Remove both ZIP and extracted directory for clean re-download
+                    os.remove(zip_file_path)
+                    import shutil
+                    shutil.rmtree(full_model_path)
+            else:
+                # Extracted directory exists but no ZIP file - assume it's valid
+                # (ZIP was cleaned up after successful extraction)
+                return full_model_path
+
+        # Download and extract the file
+        Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
+        try:
+            request.urlretrieve(url=self.path_or_url, filename=zip_file_path)
+
+            # Verify the downloaded file
+            if not verify_file_hash(zip_file_path, expected_hash):
+                if expected_hash is not None:
+                    os.remove(zip_file_path)
+                    raise RuntimeError(
+                        f"Downloaded model file from {self.path_or_url} failed integrity check. "
+                        f"This may indicate a network issue or corrupted download."
+                    )
+
+            with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
+                zip_ref.extractall(self.cache_dir)
+        finally:
+            # Clean up ZIP file after successful extraction
+            if os.path.exists(zip_file_path):
+                os.remove(zip_file_path)
+
         return full_model_path
 
     @property

etna/models/nn/timesfm.py

@@ -12,6 +12,7 @@
 import pandas as pd
 
 from etna import SETTINGS
+from etna.core.utils import get_known_hash, verify_file_hash
 from etna.datasets import TSDataset
 from etna.distributions import BaseDistribution
 from etna.models.base import NonPredictionIntervalContextRequiredAbstractModel
@@ -152,12 +153,37 @@ def _is_url(self):
         return self.path_or_url.startswith("https://") or self.path_or_url.startswith("http://")
 
     def _download_model_from_url(self) -> str:
-        """Download model from url to local cache_dir."""
+        """Download model from url to local cache_dir with integrity verification."""
         model_file = self.path_or_url.split("/")[-1]
         full_model_path = f"{self.cache_dir}/{model_file}"
-        if not os.path.exists(full_model_path):
-            Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
-            request.urlretrieve(url=self.path_or_url, filename=full_model_path)
+        expected_hash = get_known_hash(self.path_or_url)
+
+        # Check if file exists and verify integrity
+        if os.path.exists(full_model_path):
+            if verify_file_hash(full_model_path, expected_hash):
+                return full_model_path
+            else:
+                # File exists but hash doesn't match, re-download
+                if expected_hash is not None:
+                    warnings.warn(
+                        f"Local model file hash does not match expected hash. "
+                        f"This may indicate a corrupted download. Re-downloading from {self.path_or_url}"
+                    )
+                os.remove(full_model_path)
+
+        # Download the file
+        Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
+        request.urlretrieve(url=self.path_or_url, filename=full_model_path)
+
+        # Verify the downloaded file
+        if not verify_file_hash(full_model_path, expected_hash):
+            if expected_hash is not None:
+                os.remove(full_model_path)
+                raise RuntimeError(
+                    f"Downloaded model file from {self.path_or_url} failed integrity check. "
+                    f"This may indicate a network issue or corrupted download."
+                )
+
         return full_model_path
 
     @property

etna/models/utils.py

@@ -1,5 +1,8 @@
+import os
+import warnings
 from typing import Optional
 from typing import Union
+from urllib import request
 
 import pandas as pd
 

etna/transforms/embeddings/models/ts2vec.py

@@ -12,6 +12,7 @@
 import numpy as np
 
 from etna import SETTINGS
+from etna.core.utils import get_known_hash, verify_file_hash
 from etna.transforms.embeddings.models import BaseEmbeddingModel
 
 if SETTINGS.torch_required:
@@ -307,16 +308,39 @@ def load(cls, path: Optional[pathlib.Path] = None, model_name: Optional[str] = N
         if model_name is not None:
             if path is None:
                 path = _DOWNLOAD_PATH / f"{model_name}.zip"
+
+            url = f"http://etna-github-prod.cdn-tinkoff.ru/embeddings/ts2vec/{model_name}.zip"
+            expected_hash = get_known_hash(url)
+
+            # Check if file exists and verify integrity
             if os.path.exists(path):
-                warnings.warn(
-                    f"Path {path} already exists. Model {model_name} will not be downloaded. Loading existing local model."
-                )
-            else:
+                if verify_file_hash(str(path), expected_hash):
+                    # File exists and is valid (or no hash to check)
+                    pass
+                else:
+                    # File exists but hash doesn't match, re-download
+                    if expected_hash is not None:
+                        warnings.warn(
+                            f"Local model file hash does not match expected hash. "
+                            f"This may indicate a corrupted download. Re-downloading {model_name} from {url}"
+                        )
+                    os.remove(path)
+
+            # Download if file doesn't exist (or was removed due to hash mismatch)
+            if not os.path.exists(path):
                 Path(path).parent.mkdir(exist_ok=True, parents=True)
 
                 if model_name in cls.list_models():
-                    url = f"http://etna-github-prod.cdn-tinkoff.ru/embeddings/ts2vec/{model_name}.zip"
                     request.urlretrieve(url=url, filename=path)
+
+                    # Verify the downloaded file
+                    if not verify_file_hash(str(path), expected_hash):
+                        if expected_hash is not None:
+                            os.remove(path)
+                            raise RuntimeError(
+                                f"Downloaded model file {model_name} from {url} failed integrity check. "
+                                f"This may indicate a network issue or corrupted download."
+                            )
                 else:
                     raise NotImplementedError(
                         f"Model {model_name} is not available. To get list of available models use `list_models` method."