fix: integrity verification of remote scripts

Author: krishjp

src/pruna/algorithms/base/pruna_base.py

@@ -28,6 +28,7 @@
     SAVE_FUNCTIONS,
     save_pruna_model,
 )
+from pruna.engine.utils import get_fn_name
 from pruna.logging.logger import pruna_logger
 
 
@@ -365,11 +366,7 @@ def apply(self, model: Any, smash_config: SmashConfig) -> Any:
 
         # if the registered save function is None, the original saving function remains
         if self.save_fn is not None and self.save_fn != SAVE_FUNCTIONS.reapply:
-            if isinstance(self.save_fn, functools.partial):
-                fn_name = getattr(self.save_fn.func, 'name', getattr(self.save_fn.func, '__name__', str(self.save_fn.func)))
-            else:
-                fn_name = getattr(self.save_fn, 'name', getattr(self.save_fn, '__name__', str(self.save_fn)))
-            
+            fn_name = get_fn_name(self.save_fn)
             smash_config.save_fns.append(fn_name)
 
         prefix = self.algorithm_name + "_"

src/pruna/algorithms/llama_cpp.py

@@ -15,20 +15,28 @@
 from __future__ import annotations
 
 import os
-import tempfile
 import subprocess
+import tempfile
+import shutil
+import urllib.request
+import sys
 from typing import Any, Dict
 
 from ConfigSpace import Constant, OrdinalHyperparameter
 
 from pruna.algorithms.base.pruna_base import PrunaAlgorithmBase
 from pruna.algorithms.base.tags import AlgorithmTag as tags
-from pruna.config.smash_config import SmashConfigPrefixWrapper
+from pruna.config.smash_config import SmashConfig, SmashConfigPrefixWrapper
 from pruna.engine.save import SAVE_FUNCTIONS
 from pruna.engine.model_checks import is_causal_lm, is_transformers_pipeline_with_causal_lm
+from pruna.engine.utils import verify_sha256
 from pruna.logging.logger import pruna_logger
 
 
+# SHA256 hash for the pinned version (b3600) of convert_hf_to_gguf.py
+LLAMA_CPP_CONVERSION_SCRIPT_SHA256 = "f62ab712618231b3e76050f94e45dcf94567312c209b4b99bfc142229360b018"
+
+
 class LlamaCpp(PrunaAlgorithmBase):
     """
     Implement Llama.cpp as a quantizer.
@@ -128,31 +136,35 @@ def _apply(self, model: Any, smash_config: SmashConfigPrefixWrapper) -> Any:
 
         # Create a temp directory to hold HF model, f16 GGUF, and optimized GGUF
         temp_dir = tempfile.mkdtemp()
-        hf_model_dir = os.path.join(temp_dir, "hf_model")
         f16_gguf_path = os.path.join(temp_dir, "model-f16.gguf")
         quant_gguf_path = os.path.join(temp_dir, f"model-{quantization_method}.gguf")
 
         try:
-            # save HF model
-            model_to_export.save_pretrained(hf_model_dir)
-            if hasattr(smash_config, "tokenizer") and smash_config.tokenizer:
-                smash_config.tokenizer.save_pretrained(hf_model_dir)
-
-            # download the conversion script directly from llama.cpp
-            import urllib.request
-            import sys
-            script_url = "https://raw.githubusercontent.com/ggml-org/llama.cpp/b3600/convert_hf_to_gguf.py"
-            script_path = os.path.join(temp_dir, "convert_hf_to_gguf.py")
-            urllib.request.urlretrieve(script_url, script_path)
-
-            pruna_logger.info("Converting Hugging Face model to GGUF format...")
-            convert_cmd = [
-                sys.executable, script_path,
-                hf_model_dir,
-                "--outfile", f16_gguf_path,
-                "--outtype", "f16"
-            ]
-            subprocess.run(convert_cmd, check=True)
+            # Use a TemporaryDirectory for the HF model to ensure automatic cleanup
+            with tempfile.TemporaryDirectory(dir=temp_dir) as hf_model_dir:
+                model_to_export.save_pretrained(hf_model_dir)
+                if hasattr(smash_config, "tokenizer") and smash_config.tokenizer:
+                    smash_config.tokenizer.save_pretrained(hf_model_dir)
+
+                # download the conversion script directly from llama.cpp
+                script_url = "https://raw.githubusercontent.com/ggml-org/llama.cpp/b3600/convert_hf_to_gguf.py"
+                script_path = os.path.join(hf_model_dir, "convert_hf_to_gguf.py")
+                urllib.request.urlretrieve(script_url, script_path)
+
+                if not verify_sha256(script_path, LLAMA_CPP_CONVERSION_SCRIPT_SHA256):
+                    raise ValueError(
+                        f"Integrity verification failed for {script_url}. "
+                        "The downloaded script may have been tampered with or the pinned version has changed."
+                    )
+
+                pruna_logger.info("Converting Hugging Face model to GGUF format...")
+                convert_cmd = [
+                    sys.executable, script_path,
+                    hf_model_dir,
+                    "--outfile", f16_gguf_path,
+                    "--outtype", "f16"
+                ]
+                subprocess.run(convert_cmd, check=True)
 
             # quantize the GGUF model
             if quantization_method != "f16":
@@ -194,6 +206,8 @@ def _apply(self, model: Any, smash_config: SmashConfigPrefixWrapper) -> Any:
 
         except Exception as e:
             pruna_logger.error(f"Error during llama.cpp quantization: {e}")
+            if 'temp_dir' in locals() and os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir)
             raise
 
     def import_algorithm_packages(self) -> Dict[str, Any]:

src/pruna/engine/save.py

@@ -43,7 +43,7 @@
 )
 from pruna.engine.model_checks import get_helpers, is_janus_llamagen_ar
 from pruna.engine.save_artifacts import save_artifacts
-from pruna.engine.utils import determine_dtype, monkeypatch
+from pruna.engine.utils import determine_dtype, get_fn_name, monkeypatch
 from pruna.logging.logger import pruna_logger
 
 if TYPE_CHECKING:
@@ -65,11 +65,6 @@ def save_pruna_model(model: Any, model_path: str | Path, smash_config: SmashConf
         The SmashConfig object containing the save and load functions.
     """
 
-    def get_fn_name(obj):
-        if isinstance(obj, partial):
-            return get_fn_name(obj.func)
-        return getattr(obj, 'name', getattr(obj, '__name__', str(obj)))
-
     model_path = Path(model_path)
     if not model_path.exists():
         model_path.mkdir(parents=True, exist_ok=True)
@@ -495,12 +490,15 @@ def save_model_llama_cpp(model: Any, model_path: str | Path, smash_config: Smash
         gguf_file = Path(model.model_path)
         if gguf_file.exists():
             target_file = model_path / "model.gguf"
-            shutil.copy(gguf_file, target_file)
+            shutil.move(gguf_file, target_file)
+            # Cleanup the temporary directory
+            if gguf_file.parent.exists():
+                shutil.rmtree(gguf_file.parent)
             smash_config.load_fns.append(LOAD_FUNCTIONS.llama_cpp.name)
         else:
-            pruna_logger.error(f"GGUF file not found at {gguf_file}")
+            raise FileNotFoundError(f"GGUF file not found at {gguf_file}")
     else:
-        pruna_logger.error("Llama object does not have model_path attribute.")
+        raise AttributeError("Llama object does not have model_path attribute.")
 
 
 def reapply(model: Any, model_path: str | Path, smash_config: SmashConfig) -> None:

src/pruna/engine/utils.py

@@ -16,9 +16,11 @@
 
 import contextlib
 import gc
+import hashlib
 import inspect
 import json
 from contextlib import AbstractContextManager, contextmanager
+from functools import partial
 from pathlib import Path
 from typing import Any
 
@@ -38,6 +40,48 @@ def safe_memory_cleanup() -> None:
     torch.cuda.empty_cache()
 
 
+def get_fn_name(obj: Any) -> str:
+    """
+    Get the name of a function or a partial function.
+
+    Parameters
+    ----------
+    obj : Any
+        The function or partial function to get the name of.
+
+    Returns
+    -------
+    str
+        The name of the function.
+    """
+    if isinstance(obj, partial):
+        return get_fn_name(obj.func)
+    return getattr(obj, "name", getattr(obj, "__name__", str(obj)))
+
+
+def verify_sha256(file_path: str | Path, expected_hash: str) -> bool:
+    """
+    Verify the SHA256 hash of a file.
+
+    Parameters
+    ----------
+    file_path : str | Path
+        The path to the file to verify.
+    expected_hash : str
+        The expected SHA256 hash.
+
+    Returns
+    -------
+    bool
+        True if the hash matches, False otherwise.
+    """
+    sha256_hash = hashlib.sha256()
+    with Path(file_path).open("rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+    return sha256_hash.hexdigest() == expected_hash
+
+
 def load_json_config(path: str | Path, json_name: str) -> dict:
     """
     Load and parse a JSON configuration file.